From ae72c983a7cf6863df5e7281baf88b4590527b5a Mon Sep 17 00:00:00 2001 From: Hannu Krosing Date: Sun, 4 Oct 2026 11:24:37 +0000 Subject: [PATCH v4 04/12] Refactor pgbench: extract script parsing, built-in scripts, and expression evaluation to script.c/h Move the script/command data structures (MetaCommand, QueryMode, Command, ParsedScript, BuiltinScript), PRNG/distribution/hashing/permutation functions, expression evaluator, script selector, and SQL/backslash script parser out of pgbench.c into script.c and script.h. Also move the shared client/thread state types (ConnectionStateEnum, CState, TStatus, TState) into pgbench.h so both pgbench.c and script.c can access them. No function bodies, parameter lists, expressions, or comments are modified in this commit. Why function/variable linkage changes are strictly necessary: 1. Removal of `static` on 8 functions moved to script.c: These 8 functions are defined in script.c and called from pgbench.c (and later commands.c and engine.c), requiring external linkage and `extern` prototypes in script.h: - evaluateExpr() (called by executeMetaCommand) - chooseScript() (called by advanceConnectionState) - postprocess_sql_command() (called by main) - process_file() (called by main) - process_builtin() (called by main) - listAvailableScripts() (called by main) - findBuiltin() (called by main) - parseScriptWeight() (called by main) 2. Functions whose linkage is unchanged: - syntax_error() was already non-static in master (called from exprparse.y and exprscan.l). - All 23 internal script helpers called only within script.c (getrand, getExponentialRand, getGaussianRand, computeIterativeZipfian, getZipfianRand, getHashFnv1a, getHashMurmur2, permute, isLazyFunc, evalLazyFunc, evalStandardFunc, evalFunc, getMetaCommand, parseQuery, skip_sql_comments, create_sql_command, free_command, process_backslash_command, ConditionError, CheckConditional, ParseScript, read_file_contents, addScript) retain their `static` linkage in script.c. - Note that keeping TState visible to script.h allows chooseScript() to retain its exact `chooseScript(TState *thread)` signature and body in this commit (decoupling chooseScript from TState is deferred to commit 0011). 3. Removal of `static` on 5 script globals defined in script.c: `querymode`, `QUERYMODE`, `sql_script`, `num_scripts`, and `total_weight` are defined in script.c and accessed from pgbench.c (in main, sendCommand, prepareCommand, printResults, etc.), so they are declared `extern` in script.h. `builtin_script` and `pgbench_callbacks` are used only inside script.c and remain `static`. --- src/bin/pgbench/Makefile | 1 + src/bin/pgbench/meson.build | 1 + src/bin/pgbench/pgbench.c | 6419 +++++++++++------------------------ src/bin/pgbench/pgbench.h | 253 +- src/bin/pgbench/script.c | 1946 +++++++++++ src/bin/pgbench/script.h | 170 + 6 files changed, 4434 insertions(+), 4356 deletions(-) create mode 100644 src/bin/pgbench/script.c create mode 100644 src/bin/pgbench/script.h diff --git a/src/bin/pgbench/Makefile b/src/bin/pgbench/Makefile index 1fd343e65b7..dc2917d7182 100644 --- a/src/bin/pgbench/Makefile +++ b/src/bin/pgbench/Makefile @@ -12,6 +12,7 @@ OBJS = \ exprparse.o \ exprscan.o \ poller.o \ + script.o \ stats.o \ variable.o \ pgbench.o diff --git a/src/bin/pgbench/meson.build b/src/bin/pgbench/meson.build index 770d580671f..b22df9c5563 100644 --- a/src/bin/pgbench/meson.build +++ b/src/bin/pgbench/meson.build @@ -3,6 +3,7 @@ pgbench_sources = files( 'pgbench.c', 'poller.c', + 'script.c', 'stats.c', 'variable.c', ) diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c index 0b6a42d60e5..d89ff38e077 100644 --- a/src/bin/pgbench/pgbench.c +++ b/src/bin/pgbench/pgbench.c @@ -57,70 +57,15 @@ #include "libpq-fe.h" #include "pgbench.h" #include "poller.h" +#include "script.h" #include "stats.h" #include "port/pg_bitutils.h" -/* X/Open (XSI) requires to provide M_PI, but core POSIX does not */ -#ifndef M_PI -#define M_PI 3.14159265358979323846 -#endif #define ERRCODE_T_R_SERIALIZATION_FAILURE "40001" #define ERRCODE_T_R_DEADLOCK_DETECTED "40P01" #define ERRCODE_UNDEFINED_TABLE "42P01" -/* - * Hashing constants - */ -#define FNV_PRIME UINT64CONST(0x100000001b3) -#define FNV_OFFSET_BASIS UINT64CONST(0xcbf29ce484222325) -#define MM2_MUL UINT64CONST(0xc6a4a7935bd1e995) -#define MM2_MUL_TIMES_8 UINT64CONST(0x35253c9ade8f4ca8) -#define MM2_ROT 47 - - -/* - * Multi-platform thread implementations - */ - -#ifdef WIN32 -/* Use Windows threads */ -#include -#define GETERRNO() (_dosmaperr(GetLastError()), errno) -#define THREAD_T HANDLE -#define THREAD_FUNC_RETURN_TYPE unsigned -#define THREAD_FUNC_RETURN return 0 -#define THREAD_FUNC_CC __stdcall -#define THREAD_CREATE(handle, function, arg) \ - ((*(handle) = (HANDLE) _beginthreadex(NULL, 0, (function), (arg), 0, NULL)) == 0 ? errno : 0) -#define THREAD_JOIN(handle) \ - (WaitForSingleObject(handle, INFINITE) != WAIT_OBJECT_0 ? \ - GETERRNO() : CloseHandle(handle) ? 0 : GETERRNO()) -#define THREAD_BARRIER_T SYNCHRONIZATION_BARRIER -#define THREAD_BARRIER_INIT(barrier, n) \ - (InitializeSynchronizationBarrier((barrier), (n), 0) ? 0 : GETERRNO()) -#define THREAD_BARRIER_WAIT(barrier) \ - EnterSynchronizationBarrier((barrier), \ - SYNCHRONIZATION_BARRIER_FLAGS_BLOCK_ONLY) -#define THREAD_BARRIER_DESTROY(barrier) -#else -/* Use POSIX threads */ -#include "port/pg_pthread.h" -#define THREAD_T pthread_t -#define THREAD_FUNC_RETURN_TYPE void * -#define THREAD_FUNC_RETURN return NULL -#define THREAD_FUNC_CC -#define THREAD_CREATE(handle, function, arg) \ - pthread_create((handle), NULL, (function), (arg)) -#define THREAD_JOIN(handle) \ - pthread_join((handle), NULL) -#define THREAD_BARRIER_T pthread_barrier_t -#define THREAD_BARRIER_INIT(barrier, n) \ - pthread_barrier_init((barrier), NULL, (n)) -#define THREAD_BARRIER_WAIT(barrier) pthread_barrier_wait((barrier)) -#define THREAD_BARRIER_DESTROY(barrier) pthread_barrier_destroy((barrier)) -#endif - /******************************************************************** * some configurable parameters */ @@ -131,10 +76,6 @@ #define LOG_STEP_SECONDS 5 /* seconds between log messages */ #define DEFAULT_NXACTS 10 /* default nxacts */ -#define MIN_GAUSSIAN_PARAM 2.0 /* minimum parameter for gauss */ - -#define MIN_ZIPFIAN_PARAM 1.001 /* minimum parameter for zipfian */ -#define MAX_ZIPFIAN_PARAM 1000.0 /* maximum parameter for zipfian */ static int nxacts = 0; /* number of transactions per client */ static int duration = 0; /* duration in seconds */ @@ -206,19 +147,6 @@ static int64 random_seed = -1; * end of configurable parameters */ -#define nbranches 1 /* Makes little sense to change this. Change - * -s instead */ -#define ntellers 10 -#define naccounts 100000 - -/* - * The scale factor at/beyond which 32bit integers are incapable of storing - * 64bit values. - * - * Although the actual threshold is 21474, we use 20000 because it is easier to - * document and remember, and isn't that far away from the real threshold. - */ -#define SCALE_32BIT_THRESHOLD 20000 static bool use_log; /* log transaction latencies to a file */ static bool use_quiet; /* quiet logging onto stderr */ @@ -265,13 +193,11 @@ static const char *dbName = NULL; static char *logfile_prefix = NULL; static const char *progname; -#define WSEP '@' /* weight separator */ static volatile sig_atomic_t timer_exceeded = false; /* flag from signal * handler */ -#define MAX_SCRIPTS 128 /* max number of SQL scripts allowed */ #define SHELL_COMMAND_SIZE 256 /* maximum size allowed for shell command */ @@ -282,370 +208,25 @@ static volatile sig_atomic_t timer_exceeded = false; /* flag from signal static pg_time_usec_t epoch_shift; -/* - * Transaction status at the end of a command. - */ -typedef enum TStatus -{ - TSTATUS_IDLE, - TSTATUS_IN_BLOCK, - TSTATUS_CONN_ERROR, - TSTATUS_OTHER_ERROR, -} TStatus; - /* Various random sequences are initialized from this one. */ static pg_prng_state base_random_sequence; /* Synchronization barrier for start and connection */ static THREAD_BARRIER_T barrier; -/* - * Connection state machine states. - */ -typedef enum -{ - /* - * The client must first choose a script to execute. Once chosen, it can - * either be throttled (state CSTATE_PREPARE_THROTTLE under --rate), start - * right away (state CSTATE_START_TX) or not start at all if the timer was - * exceeded (state CSTATE_FINISHED). - */ - CSTATE_CHOOSE_SCRIPT, - - /* - * CSTATE_START_TX performs start-of-transaction processing. Establishes - * a new connection for the transaction in --connect mode, records the - * transaction start time, and proceed to the first command. - * - * Note: once a script is started, it will either error or run till its - * end, where it may be interrupted. It is not interrupted while running, - * so pgbench --time is to be understood as tx are allowed to start in - * that time, and will finish when their work is completed. - */ - CSTATE_START_TX, - - /* - * In CSTATE_PREPARE_THROTTLE state, we calculate when to begin the next - * transaction, and advance to CSTATE_THROTTLE. CSTATE_THROTTLE state - * sleeps until that moment, then advances to CSTATE_START_TX, or - * CSTATE_FINISHED if the next transaction would start beyond the end of - * the run. - */ - CSTATE_PREPARE_THROTTLE, - CSTATE_THROTTLE, - - /* - * We loop through these states, to process each command in the script: - * - * CSTATE_START_COMMAND starts the execution of a command. On a SQL - * command, the command is sent to the server, and we move to - * CSTATE_WAIT_RESULT state unless in pipeline mode. On a \sleep - * meta-command, the timer is set, and we enter the CSTATE_SLEEP state to - * wait for it to expire. Other meta-commands are executed immediately. If - * the command about to start is actually beyond the end of the script, - * advance to CSTATE_END_TX. - * - * CSTATE_WAIT_RESULT waits until we get a result set back from the server - * for the current command. - * - * CSTATE_SLEEP waits until the end of \sleep. - * - * CSTATE_END_COMMAND records the end-of-command timestamp, increments the - * command counter, and loops back to CSTATE_START_COMMAND state. - * - * CSTATE_SKIP_COMMAND is used by conditional branches which are not - * executed. It quickly skip commands that do not need any evaluation. - * This state can move forward several commands, till there is something - * to do or the end of the script. - */ - CSTATE_START_COMMAND, - CSTATE_WAIT_RESULT, - CSTATE_SLEEP, - CSTATE_END_COMMAND, - CSTATE_SKIP_COMMAND, - - /* - * States for failed commands. - * - * If the SQL/meta command fails, in CSTATE_ERROR clean up after an error: - * (1) clear the conditional stack; (2) if we have an unterminated - * (possibly failed) transaction block, send the rollback command to the - * server and wait for the result in CSTATE_WAIT_ROLLBACK_RESULT. If - * something goes wrong with rolling back, go to CSTATE_ABORTED. - * - * But if everything is ok we are ready for future transactions: if this - * is a serialization or deadlock error and we can re-execute the - * transaction from the very beginning, go to CSTATE_RETRY; otherwise go - * to CSTATE_FAILURE. - * - * In CSTATE_RETRY report an error, set the same parameters for the - * transaction execution as in the previous tries and process the first - * transaction command in CSTATE_START_COMMAND. - * - * In CSTATE_FAILURE report a failure, set the parameters for the - * transaction execution as they were before the first run of this - * transaction (except for a random state) and go to CSTATE_END_TX to - * complete this transaction. - */ - CSTATE_ERROR, - CSTATE_WAIT_ROLLBACK_RESULT, - CSTATE_RETRY, - CSTATE_FAILURE, - - /* - * CSTATE_END_TX performs end-of-transaction processing. It calculates - * latency, and logs the transaction. In --connect mode, it closes the - * current connection. - * - * Then either starts over in CSTATE_CHOOSE_SCRIPT, or enters - * CSTATE_FINISHED if we have no more work to do. - */ - CSTATE_END_TX, - - /* - * Final states. CSTATE_ABORTED means that the script execution was - * aborted because a command failed, CSTATE_FINISHED means success. - */ - CSTATE_ABORTED, - CSTATE_FINISHED, -} ConnectionStateEnum; - -/* - * Connection state. - */ -typedef struct -{ - PGconn *con; /* connection handle to DB */ - int id; /* client No. */ - ConnectionStateEnum state; /* state machine's current state. */ - ConditionalStack cstack; /* enclosing conditionals state */ - - /* - * Separate randomness for each client. This is used for random functions - * PGBENCH_RANDOM_* during the execution of the script. - */ - pg_prng_state cs_func_rs; - - int use_file; /* index in sql_script for this client */ - int command; /* command number in script */ - int num_syncs; /* number of ongoing sync commands */ - - /* client variables */ - Variables variables; - - /* various times about current transaction in microseconds */ - pg_time_usec_t txn_scheduled; /* scheduled start time of transaction */ - pg_time_usec_t sleep_until; /* scheduled start time of next cmd */ - pg_time_usec_t txn_begin; /* used for measuring schedule lag times */ - pg_time_usec_t stmt_begin; /* used for measuring statement latencies */ - - /* whether client prepared each command of each script */ - bool **prepared; - - /* - * For processing failures and repeating transactions with serialization - * or deadlock errors: - */ - EStatus estatus; /* the error status of the current transaction - * execution; this is ESTATUS_NO_ERROR if - * there were no errors */ - pg_prng_state random_state; /* random state */ - uint32 tries; /* how many times have we already tried the - * current transaction? */ - - /* per client collected stats */ - int64 cnt; /* client transaction count, for -t; skipped - * and failed transactions are also counted - * here */ -} CState; - -/* - * Thread state - */ -typedef struct -{ - int tid; /* thread id */ - THREAD_T thread; /* thread handle */ - CState *state; /* array of CState */ - int nstate; /* length of state[] */ - - /* - * Separate randomness for each thread. Each thread option uses its own - * random state to make all of them independent of each other and - * therefore deterministic at the thread level. - */ - pg_prng_state ts_choose_rs; /* random state for selecting a script */ - pg_prng_state ts_throttle_rs; /* random state for transaction throttling */ - pg_prng_state ts_sample_rs; /* random state for log sampling */ - - int64 throttle_trigger; /* previous/next throttling (us) */ - FILE *logfile; /* where to log, or NULL */ - - /* per thread collected stats in microseconds */ - pg_time_usec_t create_time; /* thread creation time */ - pg_time_usec_t started_time; /* thread is running */ - pg_time_usec_t bench_start; /* thread is benchmarking */ - pg_time_usec_t conn_duration; /* cumulated connection and disconnection - * delays */ - - StatsData stats; - int64 latency_late; /* count executed but late transactions */ -} TState; - -/* - * queries read from files - */ -#define SQL_COMMAND 1 -#define META_COMMAND 2 - -/* - * max number of backslash command arguments or SQL variables, - * including the command or SQL statement itself - */ -#define MAX_ARGS 256 - -typedef enum MetaCommand -{ - META_NONE, /* not a known meta-command */ - META_SET, /* \set */ - META_SETSHELL, /* \setshell */ - META_SHELL, /* \shell */ - META_SLEEP, /* \sleep */ - META_GSET, /* \gset */ - META_ASET, /* \aset */ - META_IF, /* \if */ - META_ELIF, /* \elif */ - META_ELSE, /* \else */ - META_ENDIF, /* \endif */ - META_STARTPIPELINE, /* \startpipeline */ - META_SYNCPIPELINE, /* \syncpipeline */ - META_ENDPIPELINE, /* \endpipeline */ -} MetaCommand; - -typedef enum QueryMode -{ - QUERY_SIMPLE, /* simple query */ - QUERY_EXTENDED, /* extended query */ - QUERY_PREPARED, /* extended query with prepared statements */ - NUM_QUERYMODE -} QueryMode; - -static QueryMode querymode = QUERY_SIMPLE; -static const char *const QUERYMODE[] = {"simple", "extended", "prepared"}; - -/* - * struct Command represents one command in a script. - * - * lines The raw, possibly multi-line command text. Variable substitution - * not applied. - * first_line A short, single-line extract of 'lines', for error reporting. - * type SQL_COMMAND or META_COMMAND - * meta The type of meta-command, with META_NONE/GSET/ASET if command - * is SQL. - * argc Number of arguments of the command, 0 if not yet processed. - * argv Command arguments, the first of which is the command or SQL - * string itself. For SQL commands, after post-processing - * argv[0] is the same as 'lines' with variables substituted. - * prepname The name that this command is prepared under, in prepare mode - * varprefix SQL commands terminated with \gset or \aset have this set - * to a non NULL value. If nonempty, it's used to prefix the - * variable name that receives the value. - * aset do gset on all possible queries of a combined query (\;). - * expr Parsed expression, if needed. - * stats Time spent in this command. - * retries Number of retries after a serialization or deadlock error in the - * current command. - * failures Number of errors in the current command that were not retried. - */ -typedef struct Command -{ - PQExpBufferData lines; - char *first_line; - int type; - MetaCommand meta; - int argc; - char *argv[MAX_ARGS]; - char *prepname; - char *varprefix; - PgBenchExpr *expr; - SimpleStats stats; - int64 retries; - int64 failures; -} Command; - -typedef struct ParsedScript -{ - const char *desc; /* script descriptor (eg, file name) */ - int weight; /* selection weight */ - Command **commands; /* NULL-terminated array of Commands */ - StatsData stats; /* total time spent in script */ -} ParsedScript; - -static ParsedScript sql_script[MAX_SCRIPTS]; /* SQL script files */ -static int num_scripts; /* number of scripts in sql_script[] */ -static int64 total_weight = 0; static bool verbose_errors = false; /* print verbose messages of all errors */ static bool exit_on_abort = false; /* exit when any client is aborted */ static bool continue_on_error = false; /* continue after errors */ -/* Builtin test scripts */ -typedef struct BuiltinScript -{ - const char *name; /* very short name for -b ... */ - const char *desc; /* short description */ - const char *script; /* actual pgbench script */ -} BuiltinScript; - -static const BuiltinScript builtin_script[] = -{ - { - "tpcb-like", - "", - "\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n" - "\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n" - "\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n" - "\\set delta random(-5000, 5000)\n" - "BEGIN;\n" - "UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;\n" - "SELECT abalance FROM pgbench_accounts WHERE aid = :aid;\n" - "UPDATE pgbench_tellers SET tbalance = tbalance + :delta WHERE tid = :tid;\n" - "UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid;\n" - "INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);\n" - "END;\n" - }, - { - "simple-update", - "", - "\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n" - "\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n" - "\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n" - "\\set delta random(-5000, 5000)\n" - "BEGIN;\n" - "UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;\n" - "SELECT abalance FROM pgbench_accounts WHERE aid = :aid;\n" - "INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);\n" - "END;\n" - }, - { - "select-only", - "", - "\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n" - "SELECT abalance FROM pgbench_accounts WHERE aid = :aid;\n" - } -}; - /* Function prototypes */ -static bool evaluateExpr(CState *st, PgBenchExpr *expr, - PgBenchValue *retval); static ConnectionStateEnum executeMetaCommand(CState *st, pg_time_usec_t *now); static void doLog(TState *thread, CState *st, StatsData *agg, bool tx_skipped, double latency, double lag); static void processXactStats(TState *thread, CState *st, pg_time_usec_t *now, bool tx_skipped, StatsData *agg); -static void addScript(const ParsedScript *script); static THREAD_FUNC_RETURN_TYPE THREAD_FUNC_CC threadRun(void *arg); static void finishCon(CState *st); static void setalarm(int seconds); @@ -653,10 +234,6 @@ static void setalarm(int seconds); /* callback used to build rows for COPY during data loading */ typedef void (*initRowMethod) (PQExpBufferData *sql, int64 curr); -/* callback functions for our flex lexer */ -static const PsqlScanCallbacks pgbench_callbacks = { - NULL, /* don't need get_variable functionality */ -}; static char get_table_relkind(PGconn *con, const char *table) @@ -778,83 +355,6 @@ initRandomState(pg_prng_state *state) } -/* - * random number generator: uniform distribution from min to max inclusive. - * - * Although the limits are expressed as int64, you can't generate the full - * int64 range in one call, because the difference of the limits mustn't - * overflow int64. This is not checked. - */ -static int64 -getrand(pg_prng_state *state, int64 min, int64 max) -{ - return min + (int64) pg_prng_uint64_range(state, 0, max - min); -} - -/* - * random number generator: exponential distribution from min to max inclusive. - * the parameter is so that the density of probability for the last cut-off max - * value is exp(-parameter). - */ -static int64 -getExponentialRand(pg_prng_state *state, int64 min, int64 max, - double parameter) -{ - double cut, - uniform, - rand; - - /* abort if wrong parameter, but must really be checked beforehand */ - Assert(parameter > 0.0); - cut = exp(-parameter); - /* pg_prng_double value in [0, 1), uniform in (0, 1] */ - uniform = 1.0 - pg_prng_double(state); - - /* - * inner expression in (cut, 1] (if parameter > 0), rand in [0, 1) - */ - Assert((1.0 - cut) != 0.0); - rand = -log(cut + (1.0 - cut) * uniform) / parameter; - /* return int64 random number within between min and max */ - return min + (int64) ((max - min + 1) * rand); -} - -/* random number generator: gaussian distribution from min to max inclusive */ -static int64 -getGaussianRand(pg_prng_state *state, int64 min, int64 max, - double parameter) -{ - double stdev; - double rand; - - /* abort if parameter is too low, but must really be checked beforehand */ - Assert(parameter >= MIN_GAUSSIAN_PARAM); - - /* - * Get normally-distributed random number in the range -parameter <= stdev - * < parameter. - * - * This loop is executed until the number is in the expected range. - * - * As the minimum parameter is 2.0, the probability of looping is low: - * sqrt(-2 ln(r)) <= 2 => r >= e^{-2} ~ 0.135, then when taking the - * average sinus multiplier as 2/pi, we have a 8.6% looping probability in - * the worst case. For a parameter value of 5.0, the looping probability - * is about e^{-5} * 2 / pi ~ 0.43%. - */ - do - { - stdev = pg_prng_double_normal(state); - } - while (stdev < -parameter || stdev >= parameter); - - /* stdev is in [-parameter, parameter), normalization to [0,1) */ - rand = (stdev + parameter) / (parameter * 2.0); - - /* return int64 random number within between min and max */ - return min + (int64) ((max - min + 1) * rand); -} - /* * random number generator: generate a value, such that the series of values * will approximate a Poisson distribution centered on the given value. @@ -877,203 +377,6 @@ getPoissonRand(pg_prng_state *state, double center) return (int64) (-log(uniform) * center + 0.5); } -/* - * Computing zipfian using rejection method, based on - * "Non-Uniform Random Variate Generation", - * Luc Devroye, p. 550-551, Springer 1986. - * - * This works for s > 1.0, but may perform badly for s very close to 1.0. - */ -static int64 -computeIterativeZipfian(pg_prng_state *state, int64 n, double s) -{ - double b = pow(2.0, s - 1.0); - double x, - t, - u, - v; - - /* Ensure n is sane */ - if (n <= 1) - return 1; - - while (true) - { - /* random variates */ - u = pg_prng_double(state); - v = pg_prng_double(state); - - x = floor(pow(u, -1.0 / (s - 1.0))); - - t = pow(1.0 + 1.0 / x, s - 1.0); - /* reject if too large or out of bound */ - if (v * x * (t - 1.0) / (b - 1.0) <= t / b && x <= n) - break; - } - return (int64) x; -} - -/* random number generator: zipfian distribution from min to max inclusive */ -static int64 -getZipfianRand(pg_prng_state *state, int64 min, int64 max, double s) -{ - int64 n = max - min + 1; - - /* abort if parameter is invalid */ - Assert(MIN_ZIPFIAN_PARAM <= s && s <= MAX_ZIPFIAN_PARAM); - - return min - 1 + computeIterativeZipfian(state, n, s); -} - -/* - * FNV-1a hash function - */ -static int64 -getHashFnv1a(int64 val, uint64 seed) -{ - int64 result; - int i; - - result = FNV_OFFSET_BASIS ^ seed; - for (i = 0; i < 8; ++i) - { - int32 octet = val & 0xff; - - val = val >> 8; - result = result ^ octet; - result = result * FNV_PRIME; - } - - return result; -} - -/* - * Murmur2 hash function - * - * Based on original work of Austin Appleby - * https://github.com/aappleby/smhasher/blob/master/src/MurmurHash2.cpp - */ -static int64 -getHashMurmur2(int64 val, uint64 seed) -{ - uint64 result = seed ^ MM2_MUL_TIMES_8; /* sizeof(int64) */ - uint64 k = (uint64) val; - - k *= MM2_MUL; - k ^= k >> MM2_ROT; - k *= MM2_MUL; - - result ^= k; - result *= MM2_MUL; - - result ^= result >> MM2_ROT; - result *= MM2_MUL; - result ^= result >> MM2_ROT; - - return (int64) result; -} - -/* - * Pseudorandom permutation function - * - * For small sizes, this generates each of the (size!) possible permutations - * of integers in the range [0, size) with roughly equal probability. Once - * the size is larger than 20, the number of possible permutations exceeds the - * number of distinct states of the internal pseudorandom number generator, - * and so not all possible permutations can be generated, but the permutations - * chosen should continue to give the appearance of being random. - * - * THIS FUNCTION IS NOT CRYPTOGRAPHICALLY SECURE. - * DO NOT USE FOR SUCH PURPOSE. - */ -static int64 -permute(const int64 val, const int64 isize, const int64 seed) -{ - /* using a high-end PRNG is probably overkill */ - pg_prng_state state; - uint64 size; - uint64 v; - int masklen; - uint64 mask; - int i; - - if (isize < 2) - return 0; /* nothing to permute */ - - /* Initialize prng state using the seed */ - pg_prng_seed(&state, (uint64) seed); - - /* Computations are performed on unsigned values */ - size = (uint64) isize; - v = (uint64) val % size; - - /* Mask to work modulo largest power of 2 less than or equal to size */ - masklen = pg_leftmost_one_pos64(size); - mask = (((uint64) 1) << masklen) - 1; - - /* - * Permute the input value by applying several rounds of pseudorandom - * bijective transformations. The intention here is to distribute each - * input uniformly randomly across the range, and separate adjacent inputs - * approximately uniformly randomly from each other, leading to a fairly - * random overall choice of permutation. - * - * To separate adjacent inputs, we multiply by a random number modulo - * (mask + 1), which is a power of 2. For this to be a bijection, the - * multiplier must be odd. Since this is known to lead to less randomness - * in the lower bits, we also apply a rotation that shifts the topmost bit - * into the least significant bit. In the special cases where size <= 3, - * mask = 1 and each of these operations is actually a no-op, so we also - * XOR the value with a different random number to inject additional - * randomness. Since the size is generally not a power of 2, we apply - * this bijection on overlapping upper and lower halves of the input. - * - * To distribute the inputs uniformly across the range, we then also apply - * a random offset modulo the full range. - * - * Taken together, these operations resemble a modified linear - * congruential generator, as is commonly used in pseudorandom number - * generators. The number of rounds is fairly arbitrary, but six has been - * found empirically to give a fairly good tradeoff between performance - * and uniform randomness. For small sizes it selects each of the (size!) - * possible permutations with roughly equal probability. For larger - * sizes, not all permutations can be generated, but the intended random - * spread is still produced. - */ - for (i = 0; i < 6; i++) - { - uint64 m, - r, - t; - - /* Random multiply (by an odd number), XOR and rotate of lower half */ - m = (pg_prng_uint64(&state) & mask) | 1; - r = pg_prng_uint64(&state) & mask; - if (v <= mask) - { - v = ((v * m) ^ r) & mask; - v = ((v << 1) & mask) | (v >> (masklen - 1)); - } - - /* Random multiply (by an odd number), XOR and rotate of upper half */ - m = (pg_prng_uint64(&state) & mask) | 1; - r = pg_prng_uint64(&state) & mask; - t = size - 1 - v; - if (t <= mask) - { - t = ((t * m) ^ r) & mask; - t = ((t << 1) & mask) | (t >> (masklen - 1)); - v = size - 1 - t; - } - - /* Random offset */ - r = pg_prng_uint64_range(&state, 0, size - 1); - v = (v + r) % size; - } - - return (int64) v; -} - /* call PQexec() and exit() on failure */ static void @@ -1183,4161 +486,2569 @@ getQueryParams(Variables *variables, const Command *command, } +/* + * Run a shell command. The result is assigned to the variable if not NULL. + * Return true if succeeded, or false on error. + */ static bool -isLazyFunc(PgBenchFunction func) -{ - return func == PGBENCH_AND || func == PGBENCH_OR || func == PGBENCH_CASE; -} - -/* lazy evaluation of some functions */ -static bool -evalLazyFunc(CState *st, - PgBenchFunction func, PgBenchExprLink *args, PgBenchValue *retval) +runShellCommand(Variables *variables, char *variable, char **argv, int argc) { - PgBenchValue a1, - a2; - bool ba1, - ba2; - - Assert(isLazyFunc(func) && args != NULL && args->next != NULL); - - /* args points to first condition */ - if (!evaluateExpr(st, args->expr, &a1)) - return false; - - /* second condition for AND/OR and corresponding branch for CASE */ - args = args->next; + char command[SHELL_COMMAND_SIZE]; + int i, + len = 0; + FILE *fp; + char res[64]; + char *endptr; + int retval; - switch (func) + /*---------- + * Join arguments with whitespace separators. Arguments starting with + * exactly one colon are treated as variables: + * name - append a string "name" + * :var - append a variable named 'var' + * ::name - append a string ":name" + *---------- + */ + for (i = 0; i < argc; i++) { - case PGBENCH_AND: - if (a1.type == PGBT_NULL) - { - setNullValue(retval); - return true; - } + char *arg; + int arglen; - if (!coerceToBool(&a1, &ba1)) - return false; + if (argv[i][0] != ':') + { + arg = argv[i]; /* a string literal */ + } + else if (argv[i][1] == ':') + { + arg = argv[i] + 1; /* a string literal starting with colons */ + } + else if ((arg = getVariable(variables, argv[i] + 1)) == NULL) + { + pg_log_error("%s: undefined variable \"%s\"", argv[0], argv[i]); + return false; + } - if (!ba1) - { - setBoolValue(retval, false); - return true; - } - - if (!evaluateExpr(st, args->expr, &a2)) - return false; - - if (a2.type == PGBT_NULL) - { - setNullValue(retval); - return true; - } - else if (!coerceToBool(&a2, &ba2)) - return false; - else - { - setBoolValue(retval, ba2); - return true; - } - - return true; - - case PGBENCH_OR: - - if (a1.type == PGBT_NULL) - { - setNullValue(retval); - return true; - } - - if (!coerceToBool(&a1, &ba1)) - return false; - - if (ba1) - { - setBoolValue(retval, true); - return true; - } - - if (!evaluateExpr(st, args->expr, &a2)) - return false; - - if (a2.type == PGBT_NULL) - { - setNullValue(retval); - return true; - } - else if (!coerceToBool(&a2, &ba2)) - return false; - else - { - setBoolValue(retval, ba2); - return true; - } - - case PGBENCH_CASE: - /* when true, execute branch */ - if (valueTruth(&a1)) - return evaluateExpr(st, args->expr, retval); - - /* now args contains next condition or final else expression */ - args = args->next; - - /* final else case? */ - if (args->next == NULL) - return evaluateExpr(st, args->expr, retval); - - /* no, another when, proceed */ - return evalLazyFunc(st, PGBENCH_CASE, args, retval); + arglen = strlen(arg); + if (len + arglen + (i > 0 ? 1 : 0) >= SHELL_COMMAND_SIZE - 1) + { + pg_log_error("%s: shell command is too long", argv[0]); + return false; + } - default: - /* internal error, cannot get here */ - Assert(0); - break; + if (i > 0) + command[len++] = ' '; + memcpy(command + len, arg, arglen); + len += arglen; } - return false; -} -/* maximum number of function arguments */ -#define MAX_FARGS 16 + command[len] = '\0'; -/* - * Recursive evaluation of standard functions, - * which do not require lazy evaluation. - */ -static bool -evalStandardFunc(CState *st, - PgBenchFunction func, PgBenchExprLink *args, - PgBenchValue *retval) -{ - /* evaluate all function arguments */ - int nargs = 0; - PgBenchValue vargs[MAX_FARGS] = {0}; - PgBenchExprLink *l = args; - bool has_null = false; + fflush(NULL); /* needed before either system() or popen() */ - for (nargs = 0; nargs < MAX_FARGS && l != NULL; nargs++, l = l->next) + /* Fast path for non-assignment case */ + if (variable == NULL) { - if (!evaluateExpr(st, l->expr, &vargs[nargs])) + if (system(command)) + { + if (!timer_exceeded) + pg_log_error("%s: could not launch shell command", argv[0]); return false; - has_null |= vargs[nargs].type == PGBT_NULL; + } + return true; } - if (l != NULL) + /* Execute the command with pipe and read the standard output. */ + if ((fp = popen(command, "r")) == NULL) { - pg_log_error("too many function arguments, maximum is %d", MAX_FARGS); + pg_log_error("%s: could not launch shell command", argv[0]); return false; } - - /* NULL arguments */ - if (has_null && func != PGBENCH_IS && func != PGBENCH_DEBUG) + if (fgets(res, sizeof(res), fp) == NULL) { - setNullValue(retval); - return true; + if (!timer_exceeded) + pg_log_error("%s: could not read result of shell command", argv[0]); + (void) pclose(fp); + return false; } - - /* then evaluate function */ - switch (func) + if (pclose(fp) < 0) { - /* overloaded operators */ - case PGBENCH_ADD: - case PGBENCH_SUB: - case PGBENCH_MUL: - case PGBENCH_DIV: - case PGBENCH_MOD: - case PGBENCH_EQ: - case PGBENCH_NE: - case PGBENCH_LE: - case PGBENCH_LT: - { - PgBenchValue *lval = &vargs[0], - *rval = &vargs[1]; - - Assert(nargs == 2); + pg_log_error("%s: could not run shell command: %m", argv[0]); + return false; + } - /* overloaded type management, double if some double */ - if ((lval->type == PGBT_DOUBLE || - rval->type == PGBT_DOUBLE) && func != PGBENCH_MOD) - { - double ld, - rd; + /* Check whether the result is an integer and assign it to the variable */ + retval = (int) strtol(res, &endptr, 10); + while (*endptr != '\0' && isspace((unsigned char) *endptr)) + endptr++; + if (*res == '\0' || *endptr != '\0') + { + pg_log_error("%s: shell command must return an integer (not \"%s\")", argv[0], res); + return false; + } + if (!putVariableInt(variables, "setshell", variable, retval)) + return false; - if (!coerceToDouble(lval, &ld) || - !coerceToDouble(rval, &rd)) - return false; + pg_log_debug("%s: shell parameter name: \"%s\", value: \"%s\"", argv[0], argv[1], res); - switch (func) - { - case PGBENCH_ADD: - setDoubleValue(retval, ld + rd); - return true; + return true; +} - case PGBENCH_SUB: - setDoubleValue(retval, ld - rd); - return true; +/* + * Report the abortion of the client when processing SQL commands. + */ +static void +commandFailed(CState *st, const char *cmd, const char *message) +{ + pg_log_error("client %d aborted in command %d (%s) of script %d; %s", + st->id, st->command, cmd, st->use_file, message); +} - case PGBENCH_MUL: - setDoubleValue(retval, ld * rd); - return true; +/* + * Report the error in the command while the script is executing. + */ +static void +commandError(CState *st, const char *message) +{ + /* + * Errors should only be detected during an SQL command or the + * \endpipeline meta command. Any other case triggers an assertion + * failure. + */ + Assert(sql_script[st->use_file].commands[st->command]->type == SQL_COMMAND || + sql_script[st->use_file].commands[st->command]->meta == META_ENDPIPELINE); - case PGBENCH_DIV: - setDoubleValue(retval, ld / rd); - return true; + pg_log_info("client %d got an error in command %d (SQL) of script %d; %s", + st->id, st->command, st->use_file, message); +} - case PGBENCH_EQ: - setBoolValue(retval, ld == rd); - return true; - case PGBENCH_NE: - setBoolValue(retval, ld != rd); - return true; +/* + * Allocate space for CState->prepared: we need one boolean for each command + * of each script. + */ +static void +allocCStatePrepared(CState *st) +{ + Assert(st->prepared == NULL); - case PGBENCH_LE: - setBoolValue(retval, ld <= rd); - return true; + st->prepared = pg_malloc_array(bool *, num_scripts); + for (int i = 0; i < num_scripts; i++) + { + ParsedScript *script = &sql_script[i]; + int numcmds; - case PGBENCH_LT: - setBoolValue(retval, ld < rd); - return true; + for (numcmds = 0; script->commands[numcmds] != NULL; numcmds++) + ; + st->prepared[i] = pg_malloc0_array(bool, numcmds); + } +} - default: - /* cannot get here */ - Assert(0); - } - } - else /* we have integer operands, or % */ - { - int64 li, - ri, - res; +/* + * Prepare the SQL command from st->use_file at command_num. + */ +static void +prepareCommand(CState *st, int command_num) +{ + Command *command = sql_script[st->use_file].commands[command_num]; - if (!coerceToInt(lval, &li) || - !coerceToInt(rval, &ri)) - return false; + /* No prepare for non-SQL commands */ + if (command->type != SQL_COMMAND) + return; - switch (func) - { - case PGBENCH_ADD: - if (pg_add_s64_overflow(li, ri, &res)) - { - pg_log_error("bigint add out of range"); - return false; - } - setIntValue(retval, res); - return true; - - case PGBENCH_SUB: - if (pg_sub_s64_overflow(li, ri, &res)) - { - pg_log_error("bigint sub out of range"); - return false; - } - setIntValue(retval, res); - return true; - - case PGBENCH_MUL: - if (pg_mul_s64_overflow(li, ri, &res)) - { - pg_log_error("bigint mul out of range"); - return false; - } - setIntValue(retval, res); - return true; - - case PGBENCH_EQ: - setBoolValue(retval, li == ri); - return true; - - case PGBENCH_NE: - setBoolValue(retval, li != ri); - return true; - - case PGBENCH_LE: - setBoolValue(retval, li <= ri); - return true; - - case PGBENCH_LT: - setBoolValue(retval, li < ri); - return true; - - case PGBENCH_DIV: - case PGBENCH_MOD: - if (ri == 0) - { - pg_log_error("division by zero"); - return false; - } - /* special handling of -1 divisor */ - if (ri == -1) - { - if (func == PGBENCH_DIV) - { - /* overflow check (needed for INT64_MIN) */ - if (li == PG_INT64_MIN) - { - pg_log_error("bigint div out of range"); - return false; - } - else - setIntValue(retval, -li); - } - else - setIntValue(retval, 0); - return true; - } - /* else divisor is not -1 */ - if (func == PGBENCH_DIV) - setIntValue(retval, li / ri); - else /* func == PGBENCH_MOD */ - setIntValue(retval, li % ri); - - return true; + if (!st->prepared) + allocCStatePrepared(st); - default: - /* cannot get here */ - Assert(0); - } - } + if (!st->prepared[st->use_file][command_num]) + { + PGresult *res; - Assert(0); - return false; /* NOTREACHED */ - } + pg_log_debug("client %d preparing %s", st->id, command->prepname); + res = PQprepare(st->con, command->prepname, + command->argv[0], command->argc - 1, NULL); + if (PQresultStatus(res) != PGRES_COMMAND_OK) + pg_log_error("%s", PQerrorMessage(st->con)); + PQclear(res); + st->prepared[st->use_file][command_num] = true; + } +} - /* integer bitwise operators */ - case PGBENCH_BITAND: - case PGBENCH_BITOR: - case PGBENCH_BITXOR: - case PGBENCH_LSHIFT: - case PGBENCH_RSHIFT: - { - int64 li, - ri; - - if (!coerceToInt(&vargs[0], &li) || !coerceToInt(&vargs[1], &ri)) - return false; - - if (func == PGBENCH_BITAND) - setIntValue(retval, li & ri); - else if (func == PGBENCH_BITOR) - setIntValue(retval, li | ri); - else if (func == PGBENCH_BITXOR) - setIntValue(retval, li ^ ri); - else if (func == PGBENCH_LSHIFT) - setIntValue(retval, li << ri); - else if (func == PGBENCH_RSHIFT) - setIntValue(retval, li >> ri); - else /* cannot get here */ - Assert(0); - - return true; - } +/* + * Prepare all the commands in the script that come after the \startpipeline + * that's at position st->command, and the first \endpipeline we find. + * + * This sets the ->prepared flag for each relevant command as well as the + * \startpipeline itself, but doesn't move the st->command counter. + */ +static void +prepareCommandsInPipeline(CState *st) +{ + int j; + Command **commands = sql_script[st->use_file].commands; - /* logical operators */ - case PGBENCH_NOT: - { - bool b; + Assert(commands[st->command]->type == META_COMMAND && + commands[st->command]->meta == META_STARTPIPELINE); - if (!coerceToBool(&vargs[0], &b)) - return false; + if (!st->prepared) + allocCStatePrepared(st); - setBoolValue(retval, !b); - return true; - } + /* + * We set the 'prepared' flag on the \startpipeline itself to flag that we + * don't need to do this next time without calling prepareCommand(), even + * though we don't actually prepare this command. + */ + if (st->prepared[st->use_file][st->command]) + return; - /* no arguments */ - case PGBENCH_PI: - setDoubleValue(retval, M_PI); - return true; + for (j = st->command + 1; commands[j] != NULL; j++) + { + if (commands[j]->type == META_COMMAND && + commands[j]->meta == META_ENDPIPELINE) + break; - /* 1 overloaded argument */ - case PGBENCH_ABS: - { - PgBenchValue *varg = &vargs[0]; + prepareCommand(st, j); + } - Assert(nargs == 1); + st->prepared[st->use_file][st->command] = true; +} - if (varg->type == PGBT_INT) - { - int64 i = varg->u.ival; +/* Send a SQL command, using the chosen querymode */ +static bool +sendCommand(CState *st, Command *command) +{ + int r; - setIntValue(retval, i < 0 ? -i : i); - } - else - { - double d = varg->u.dval; + if (querymode == QUERY_SIMPLE) + { + char *sql; - Assert(varg->type == PGBT_DOUBLE); - setDoubleValue(retval, d < 0.0 ? -d : d); - } + sql = pg_strdup(command->argv[0]); + sql = assignVariables(&st->variables, sql); - return true; - } + pg_log_debug("client %d sending %s", st->id, sql); + r = PQsendQuery(st->con, sql); + pg_free(sql); + } + else if (querymode == QUERY_EXTENDED) + { + const char *sql = command->argv[0]; + const char *params[MAX_ARGS]; - case PGBENCH_DEBUG: - { - PgBenchValue *varg = &vargs[0]; + getQueryParams(&st->variables, command, params); - Assert(nargs == 1); + pg_log_debug("client %d sending %s", st->id, sql); + r = PQsendQueryParams(st->con, sql, command->argc - 1, + NULL, params, NULL, NULL, 0); + } + else if (querymode == QUERY_PREPARED) + { + const char *params[MAX_ARGS]; - fprintf(stderr, "debug(script=%d,command=%d): ", - st->use_file, st->command + 1); + prepareCommand(st, st->command); + getQueryParams(&st->variables, command, params); - if (varg->type == PGBT_NULL) - fprintf(stderr, "null\n"); - else if (varg->type == PGBT_BOOLEAN) - fprintf(stderr, "boolean %s\n", varg->u.bval ? "true" : "false"); - else if (varg->type == PGBT_INT) - fprintf(stderr, "int " INT64_FORMAT "\n", varg->u.ival); - else if (varg->type == PGBT_DOUBLE) - fprintf(stderr, "double %.*g\n", DBL_DIG, varg->u.dval); - else /* internal error, unexpected type */ - Assert(0); + pg_log_debug("client %d sending %s", st->id, command->prepname); + r = PQsendQueryPrepared(st->con, command->prepname, command->argc - 1, + params, NULL, NULL, 0); + } + else /* unknown sql mode */ + r = 0; - *retval = *varg; + if (r == 0) + { + pg_log_debug("client %d could not send %s", st->id, command->argv[0]); + return false; + } + else + return true; +} - return true; - } +/* + * Read and discard all available results from the connection. + */ +static void +discardAvailableResults(CState *st) +{ + PGresult *res = NULL; - /* 1 double argument */ - case PGBENCH_DOUBLE: - case PGBENCH_SQRT: - case PGBENCH_LN: - case PGBENCH_EXP: - { - double dval; + for (;;) + { + res = PQgetResult(st->con); + + /* + * Read and discard results until PQgetResult() returns NULL (no more + * results) or a connection failure is detected. If the pipeline + * status is PQ_PIPELINE_ABORTED, more results may still be available + * even after PQgetResult() returns NULL, so continue reading in that + * case. + */ + if ((res == NULL && PQpipelineStatus(st->con) != PQ_PIPELINE_ABORTED) || + PQstatus(st->con) == CONNECTION_BAD) + break; + + PQclear(res); + } + PQclear(res); +} - Assert(nargs == 1); +/* + * Determine the error status based on the connection status and error code. + */ +static EStatus +getSQLErrorStatus(CState *st, const char *sqlState) +{ + discardAvailableResults(st); + if (PQstatus(st->con) == CONNECTION_BAD) + return ESTATUS_CONN_ERROR; - if (!coerceToDouble(&vargs[0], &dval)) - return false; + if (sqlState != NULL) + { + if (strcmp(sqlState, ERRCODE_T_R_SERIALIZATION_FAILURE) == 0) + return ESTATUS_SERIALIZATION_ERROR; + else if (strcmp(sqlState, ERRCODE_T_R_DEADLOCK_DETECTED) == 0) + return ESTATUS_DEADLOCK_ERROR; + } - if (func == PGBENCH_SQRT) - dval = sqrt(dval); - else if (func == PGBENCH_LN) - dval = log(dval); - else if (func == PGBENCH_EXP) - dval = exp(dval); - /* else is cast: do nothing */ + return ESTATUS_OTHER_SQL_ERROR; +} - setDoubleValue(retval, dval); - return true; - } +/* + * Returns true if this type of error can be retried. + */ +static bool +canRetryError(EStatus estatus) +{ + return (estatus == ESTATUS_SERIALIZATION_ERROR || + estatus == ESTATUS_DEADLOCK_ERROR); +} - /* 1 int argument */ - case PGBENCH_INT: - { - int64 ival; +/* + * Returns true if --continue-on-error is specified and this error allows + * processing to continue. + */ +static bool +canContinueOnError(EStatus estatus) +{ + return (continue_on_error && + estatus == ESTATUS_OTHER_SQL_ERROR); +} - Assert(nargs == 1); +/* + * Process query response from the backend. + * + * If varprefix is not NULL, it's the variable name prefix where to store + * the results of the *last* command (META_GSET) or *all* commands + * (META_ASET). + * + * Returns true if everything is A-OK, false if any error occurs. + */ +static bool +readCommandResponse(CState *st, MetaCommand meta, char *varprefix) +{ + PGresult *res; + PGresult *next_res; + int qrynum = 0; - if (!coerceToInt(&vargs[0], &ival)) - return false; + /* + * varprefix should be set only with \gset or \aset, and \endpipeline and + * SQL commands do not need it. + */ + Assert((meta == META_NONE && varprefix == NULL) || + ((meta == META_ENDPIPELINE) && varprefix == NULL) || + ((meta == META_GSET || meta == META_ASET) && varprefix != NULL)); - setIntValue(retval, ival); - return true; - } + res = PQgetResult(st->con); - /* variable number of arguments */ - case PGBENCH_LEAST: - case PGBENCH_GREATEST: - { - bool havedouble; - int i; + while (res != NULL) + { + bool is_last; - Assert(nargs >= 1); + /* peek at the next result to know whether the current is last */ + next_res = PQgetResult(st->con); + is_last = (next_res == NULL); - /* need double result if any input is double */ - havedouble = false; - for (i = 0; i < nargs; i++) + switch (PQresultStatus(res)) + { + case PGRES_COMMAND_OK: /* non-SELECT commands */ + case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */ + if (is_last && meta == META_GSET) { - if (vargs[i].type == PGBT_DOUBLE) - { - havedouble = true; - break; - } + pg_log_error("client %d script %d command %d query %d: expected one row, got %d", + st->id, st->use_file, st->command, qrynum, 0); + st->estatus = ESTATUS_META_COMMAND_ERROR; + goto error; } - if (havedouble) + break; + + case PGRES_TUPLES_OK: + if ((is_last && meta == META_GSET) || meta == META_ASET) { - double extremum; + int ntuples = PQntuples(res); + const char *context = (meta == META_ASET) ? "aset" : "gset"; - if (!coerceToDouble(&vargs[0], &extremum)) - return false; - for (i = 1; i < nargs; i++) + if (meta == META_GSET && ntuples != 1) { - double dval; - - if (!coerceToDouble(&vargs[i], &dval)) - return false; - if (func == PGBENCH_LEAST) - extremum = Min(extremum, dval); - else - extremum = Max(extremum, dval); + /* under \gset, report the error */ + pg_log_error("client %d script %d command %d query %d: expected one row, got %d", + st->id, st->use_file, st->command, qrynum, PQntuples(res)); + st->estatus = ESTATUS_META_COMMAND_ERROR; + goto error; } - setDoubleValue(retval, extremum); - } - else - { - int64 extremum; - - if (!coerceToInt(&vargs[0], &extremum)) - return false; - for (i = 1; i < nargs; i++) + else if (meta == META_ASET && ntuples <= 0) { - int64 ival; - - if (!coerceToInt(&vargs[i], &ival)) - return false; - if (func == PGBENCH_LEAST) - extremum = Min(extremum, ival); - else - extremum = Max(extremum, ival); + /* coldly skip empty result under \aset */ + break; } - setIntValue(retval, extremum); - } - return true; - } - - /* random functions */ - case PGBENCH_RANDOM: - case PGBENCH_RANDOM_EXPONENTIAL: - case PGBENCH_RANDOM_GAUSSIAN: - case PGBENCH_RANDOM_ZIPFIAN: - { - int64 imin, - imax, - delta; - - Assert(nargs >= 2); - - if (!coerceToInt(&vargs[0], &imin) || - !coerceToInt(&vargs[1], &imax)) - return false; - - /* check random range */ - if (unlikely(imin > imax)) - { - pg_log_error("empty range given to random"); - return false; - } - else if (unlikely(pg_sub_s64_overflow(imax, imin, &delta) || - pg_add_s64_overflow(delta, 1, &delta))) - { - /* prevent int overflows in random functions */ - pg_log_error("random range is too large"); - return false; - } - - if (func == PGBENCH_RANDOM) - { - Assert(nargs == 2); - setIntValue(retval, getrand(&st->cs_func_rs, imin, imax)); - } - else /* gaussian & exponential */ - { - double param; - Assert(nargs == 3); + /* store results into variables */ + for (int fld = 0; fld < PQnfields(res); fld++) + { + char *varname = PQfname(res, fld); + bool ok; - if (!coerceToDouble(&vargs[2], ¶m)) - return false; + /* allocate varname only if necessary, freed below */ + if (*varprefix != '\0') + varname = psprintf("%s%s", varprefix, varname); - if (func == PGBENCH_RANDOM_GAUSSIAN) - { - if (param < MIN_GAUSSIAN_PARAM) + /* store last row result */ + if (PQgetisnull(res, ntuples - 1, fld)) { - pg_log_error("gaussian parameter must be at least %f (not %f)", - MIN_GAUSSIAN_PARAM, param); - return false; - } + PgBenchValue nullval; - setIntValue(retval, - getGaussianRand(&st->cs_func_rs, - imin, imax, param)); - } - else if (func == PGBENCH_RANDOM_ZIPFIAN) - { - if (param < MIN_ZIPFIAN_PARAM || param > MAX_ZIPFIAN_PARAM) - { - pg_log_error("zipfian parameter must be in range [%.3f, %.0f] (not %f)", - MIN_ZIPFIAN_PARAM, MAX_ZIPFIAN_PARAM, param); - return false; + setNullValue(&nullval); + ok = putVariableValue(&st->variables, context, + varname, &nullval); } + else + ok = putVariable(&st->variables, context, varname, + PQgetvalue(res, ntuples - 1, fld)); - setIntValue(retval, - getZipfianRand(&st->cs_func_rs, imin, imax, param)); - } - else /* exponential */ - { - if (param <= 0.0) + if (!ok) { - pg_log_error("exponential parameter must be greater than zero (not %f)", - param); - return false; + /* internal error */ + pg_log_error("client %d script %d command %d query %d: error storing into variable %s", + st->id, st->use_file, st->command, qrynum, varname); + st->estatus = ESTATUS_META_COMMAND_ERROR; + goto error; } - setIntValue(retval, - getExponentialRand(&st->cs_func_rs, - imin, imax, param)); + if (*varprefix != '\0') + pfree(varname); } } + /* otherwise the result is simply thrown away by PQclear below */ + break; - return true; - } - - case PGBENCH_POW: - { - PgBenchValue *lval = &vargs[0]; - PgBenchValue *rval = &vargs[1]; - double ld, - rd; - - Assert(nargs == 2); - - if (!coerceToDouble(lval, &ld) || - !coerceToDouble(rval, &rd)) - return false; - - setDoubleValue(retval, pow(ld, rd)); - - return true; - } + case PGRES_PIPELINE_SYNC: + pg_log_debug("client %d pipeline ending, ongoing syncs: %d", + st->id, st->num_syncs); + st->num_syncs--; + if (st->num_syncs == 0 && PQexitPipelineMode(st->con) != 1) + pg_log_error("client %d failed to exit pipeline mode: %s", st->id, + PQresultErrorMessage(res)); + break; - case PGBENCH_IS: - { - Assert(nargs == 2); + case PGRES_COPY_IN: + case PGRES_COPY_OUT: + case PGRES_COPY_BOTH: + pg_log_error("COPY is not supported in pgbench, aborting"); /* - * note: this simple implementation is more permissive than - * SQL + * We need to exit the copy state. Otherwise, PQgetResult() + * will always return an empty PGresult as an effect of + * getCopyResult(), leading to an infinite loop in the error + * cleanup done below. */ - setBoolValue(retval, - vargs[0].type == vargs[1].type && - vargs[0].u.bval == vargs[1].u.bval); - return true; - } + PQendcopy(st->con); + goto error; - /* hashing */ - case PGBENCH_HASH_FNV1A: - case PGBENCH_HASH_MURMUR2: - { - int64 val, - seed; + case PGRES_NONFATAL_ERROR: + case PGRES_FATAL_ERROR: + st->estatus = getSQLErrorStatus(st, PQresultErrorField(res, + PG_DIAG_SQLSTATE)); + if (canRetryError(st->estatus) || canContinueOnError(st->estatus)) + { + if (verbose_errors) + commandError(st, PQresultErrorMessage(res)); + goto error; + } + pg_fallthrough; - Assert(nargs == 2); + default: + /* anything else is unexpected */ + pg_log_error("client %d script %d aborted in command %d query %d: %s", + st->id, st->use_file, st->command, qrynum, + PQresultErrorMessage(res)); + goto error; + } - if (!coerceToInt(&vargs[0], &val) || - !coerceToInt(&vargs[1], &seed)) - return false; + PQclear(res); + qrynum++; + res = next_res; + } - if (func == PGBENCH_HASH_MURMUR2) - setIntValue(retval, getHashMurmur2(val, seed)); - else if (func == PGBENCH_HASH_FNV1A) - setIntValue(retval, getHashFnv1a(val, seed)); - else - /* cannot get here */ - Assert(0); + if (qrynum == 0) + { + pg_log_error("client %d command %d: no results", st->id, st->command); + return false; + } - return true; - } + return true; - case PGBENCH_PERMUTE: - { - int64 val, - size, - seed; +error: + PQclear(res); + PQclear(next_res); + discardAvailableResults(st); - Assert(nargs == 3); + return false; +} - if (!coerceToInt(&vargs[0], &val) || - !coerceToInt(&vargs[1], &size) || - !coerceToInt(&vargs[2], &seed)) - return false; +/* + * Parse the argument to a \sleep command, and return the requested amount + * of delay, in microseconds. Returns true on success, false on error. + */ +static bool +evaluateSleep(Variables *variables, int argc, char **argv, int *usecs) +{ + char *var; + int usec; - if (size <= 0) - { - pg_log_error("permute size parameter must be greater than zero"); - return false; - } + if (*argv[1] == ':') + { + if ((var = getVariable(variables, argv[1] + 1)) == NULL) + { + pg_log_error("%s: undefined variable \"%s\"", argv[0], argv[1] + 1); + return false; + } - setIntValue(retval, permute(val, size, seed)); - return true; - } + usec = atoi(var); - default: - /* cannot get here */ - Assert(0); - /* dead code to avoid a compiler warning */ + /* 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]); -/* evaluate some function */ -static bool -evalFunc(CState *st, - PgBenchFunction func, PgBenchExprLink *args, PgBenchValue *retval) -{ - if (isLazyFunc(func)) - return evalLazyFunc(st, func, args, retval); + if (argc > 2) + { + if (pg_strcasecmp(argv[2], "ms") == 0) + usec *= 1000; + else if (pg_strcasecmp(argv[2], "s") == 0) + usec *= 1000000; + } else - return evalStandardFunc(st, func, args, retval); + usec *= 1000000; + + *usecs = usec; + return true; } + /* - * Recursive evaluation of an expression in a pgbench script - * using the current state of variables. - * Returns whether the evaluation was ok, - * the value itself is returned through the retval pointer. + * Returns true if the error can be retried. */ static bool -evaluateExpr(CState *st, PgBenchExpr *expr, PgBenchValue *retval) +doRetry(CState *st, pg_time_usec_t *now) { - switch (expr->etype) - { - case ENODE_CONSTANT: - { - *retval = expr->u.constant; - return true; - } + Assert(st->estatus != ESTATUS_NO_ERROR); - case ENODE_VARIABLE: - { - Variable *var; + /* We can only retry serialization or deadlock errors. */ + if (!canRetryError(st->estatus)) + return false; - if ((var = lookupVariable(&st->variables, expr->u.variable.varname)) == NULL) - { - pg_log_error("undefined variable \"%s\"", expr->u.variable.varname); - return false; - } + /* + * We must have at least one option to limit the retrying of transactions + * that got an error. + */ + Assert(max_tries || latency_limit || duration > 0); - if (!makeVariableValue(var)) - return false; + /* + * We cannot retry the error if we have reached the maximum number of + * tries. + */ + if (max_tries && st->tries >= max_tries) + return false; - *retval = var->value; - return true; - } + /* + * We cannot retry the error if we spent too much time on this + * transaction. + */ + if (latency_limit) + { + pg_time_now_lazy(now); + if (*now - st->txn_scheduled > latency_limit) + return false; + } - case ENODE_FUNCTION: - return evalFunc(st, - expr->u.function.function, - expr->u.function.args, - retval); + /* + * We cannot retry the error if the benchmark duration is over. + */ + if (timer_exceeded) + return false; - default: - /* internal error which should never occur */ - pg_fatal("unexpected enode type in evaluation: %d", expr->etype); - } + /* OK */ + return true; } /* - * Convert command name to meta-command enum identifier + * Read and discard results until the last sync point. */ -static MetaCommand -getMetaCommand(const char *cmd) +static int +discardUntilSync(CState *st) { - MetaCommand mc; - - if (cmd == NULL) - mc = META_NONE; - else if (pg_strcasecmp(cmd, "set") == 0) - mc = META_SET; - else if (pg_strcasecmp(cmd, "setshell") == 0) - mc = META_SETSHELL; - else if (pg_strcasecmp(cmd, "shell") == 0) - mc = META_SHELL; - else if (pg_strcasecmp(cmd, "sleep") == 0) - mc = META_SLEEP; - else if (pg_strcasecmp(cmd, "if") == 0) - mc = META_IF; - else if (pg_strcasecmp(cmd, "elif") == 0) - mc = META_ELIF; - else if (pg_strcasecmp(cmd, "else") == 0) - mc = META_ELSE; - else if (pg_strcasecmp(cmd, "endif") == 0) - mc = META_ENDIF; - else if (pg_strcasecmp(cmd, "gset") == 0) - mc = META_GSET; - else if (pg_strcasecmp(cmd, "aset") == 0) - mc = META_ASET; - else if (pg_strcasecmp(cmd, "startpipeline") == 0) - mc = META_STARTPIPELINE; - else if (pg_strcasecmp(cmd, "syncpipeline") == 0) - mc = META_SYNCPIPELINE; - else if (pg_strcasecmp(cmd, "endpipeline") == 0) - mc = META_ENDPIPELINE; - else - mc = META_NONE; - return mc; -} + bool received_sync = false; -/* - * Run a shell command. The result is assigned to the variable if not NULL. - * Return true if succeeded, or false on error. - */ -static bool -runShellCommand(Variables *variables, char *variable, char **argv, int argc) -{ - char command[SHELL_COMMAND_SIZE]; - int i, - len = 0; - FILE *fp; - char res[64]; - char *endptr; - int retval; + /* + * Send a Sync message to ensure at least one PGRES_PIPELINE_SYNC is + * received and to avoid an infinite loop, since all earlier ones may have + * already been received. + */ + if (!PQpipelineSync(st->con)) + { + pg_log_error("client %d aborted: failed to send a pipeline sync", + st->id); + return 0; + } - /*---------- - * Join arguments with whitespace separators. Arguments starting with - * exactly one colon are treated as variables: - * name - append a string "name" - * :var - append a variable named 'var' - * ::name - append a string ":name" - *---------- + /* + * Continue reading results until the last sync point, i.e., until + * reaching null just after PGRES_PIPELINE_SYNC. */ - for (i = 0; i < argc; i++) + for (;;) { - char *arg; - int arglen; + PGresult *res = PQgetResult(st->con); - if (argv[i][0] != ':') - { - arg = argv[i]; /* a string literal */ - } - else if (argv[i][1] == ':') - { - arg = argv[i] + 1; /* a string literal starting with colons */ - } - else if ((arg = getVariable(variables, argv[i] + 1)) == NULL) + if (PQstatus(st->con) == CONNECTION_BAD) { - pg_log_error("%s: undefined variable \"%s\"", argv[0], argv[i]); - return false; + pg_log_error("client %d aborted while rolling back the transaction after an error; perhaps the backend died while processing", + st->id); + PQclear(res); + return 0; } - arglen = strlen(arg); - if (len + arglen + (i > 0 ? 1 : 0) >= SHELL_COMMAND_SIZE - 1) + if (PQresultStatus(res) == PGRES_PIPELINE_SYNC) + received_sync = true; + else if (received_sync && res == NULL) { - pg_log_error("%s: shell command is too long", argv[0]); - return false; + /* + * Reset ongoing sync count to 0 since all PGRES_PIPELINE_SYNC + * results have been discarded. + */ + st->num_syncs = 0; + break; } - - if (i > 0) - command[len++] = ' '; - memcpy(command + len, arg, arglen); - len += arglen; - } - - command[len] = '\0'; - - fflush(NULL); /* needed before either system() or popen() */ - - /* Fast path for non-assignment case */ - if (variable == NULL) - { - if (system(command)) + else { - if (!timer_exceeded) - pg_log_error("%s: could not launch shell command", argv[0]); - return false; + /* + * If a PGRES_PIPELINE_SYNC is followed by something other than + * PGRES_PIPELINE_SYNC or NULL, another PGRES_PIPELINE_SYNC will + * appear later. Reset received_sync to false to wait for it. + */ + received_sync = false; } - return true; + PQclear(res); } - /* Execute the command with pipe and read the standard output. */ - if ((fp = popen(command, "r")) == NULL) - { - pg_log_error("%s: could not launch shell command", argv[0]); - return false; - } - if (fgets(res, sizeof(res), fp) == NULL) - { - if (!timer_exceeded) - pg_log_error("%s: could not read result of shell command", argv[0]); - (void) pclose(fp); - return false; - } - if (pclose(fp) < 0) + /* exit pipeline */ + if (PQexitPipelineMode(st->con) != 1) { - pg_log_error("%s: could not run shell command: %m", argv[0]); - return false; + pg_log_error("client %d aborted: failed to exit pipeline mode for rolling back the failed transaction", + st->id); + return 0; } + return 1; +} - /* Check whether the result is an integer and assign it to the variable */ - retval = (int) strtol(res, &endptr, 10); - while (*endptr != '\0' && isspace((unsigned char) *endptr)) - endptr++; - if (*res == '\0' || *endptr != '\0') +/* + * Get the transaction status at the end of a command especially for + * checking if we are in a (failed) transaction block. + */ +static TStatus +getTransactionStatus(PGconn *con) +{ + PGTransactionStatusType tx_status; + + tx_status = PQtransactionStatus(con); + switch (tx_status) { - pg_log_error("%s: shell command must return an integer (not \"%s\")", argv[0], res); - return false; - } - if (!putVariableInt(variables, "setshell", variable, retval)) - return false; + case PQTRANS_IDLE: + return TSTATUS_IDLE; + case PQTRANS_INTRANS: + case PQTRANS_INERROR: + return TSTATUS_IN_BLOCK; + case PQTRANS_UNKNOWN: + /* PQTRANS_UNKNOWN is expected given a broken connection */ + if (PQstatus(con) == CONNECTION_BAD) + return TSTATUS_CONN_ERROR; + pg_fallthrough; + case PQTRANS_ACTIVE: + default: - pg_log_debug("%s: shell parameter name: \"%s\", value: \"%s\"", argv[0], argv[1], res); + /* + * We cannot find out whether we are in a transaction block or + * not. Internal error which should never occur. + */ + pg_log_error("unexpected transaction status %d", tx_status); + return TSTATUS_OTHER_ERROR; + } - return true; + /* not reached */ + Assert(false); + return TSTATUS_OTHER_ERROR; } /* - * Report the abortion of the client when processing SQL commands. + * Print verbose messages of an error */ static void -commandFailed(CState *st, const char *cmd, const char *message) +printVerboseErrorMessages(CState *st, pg_time_usec_t *now, bool is_retry) { - pg_log_error("client %d aborted in command %d (%s) of script %d; %s", - st->id, st->command, cmd, st->use_file, message); -} - -/* - * Report the error in the command while the script is executing. - */ -static void -commandError(CState *st, const char *message) -{ - /* - * Errors should only be detected during an SQL command or the - * \endpipeline meta command. Any other case triggers an assertion - * failure. - */ - Assert(sql_script[st->use_file].commands[st->command]->type == SQL_COMMAND || - sql_script[st->use_file].commands[st->command]->meta == META_ENDPIPELINE); + PQExpBufferData buf; - pg_log_info("client %d got an error in command %d (SQL) of script %d; %s", - st->id, st->command, st->use_file, message); -} + initPQExpBuffer(&buf); -/* return a script number with a weighted choice. */ -static int -chooseScript(TState *thread) -{ - int i = 0; - int64 w; + printfPQExpBuffer(&buf, "client %d ", st->id); + appendPQExpBufferStr(&buf, (is_retry ? + "repeats the transaction after the error" : + "ends the failed transaction")); + appendPQExpBuffer(&buf, " (try %u", st->tries); - if (num_scripts == 1) - return 0; + /* Print max_tries if it is not unlimited. */ + if (max_tries) + appendPQExpBuffer(&buf, "/%u", max_tries); - w = getrand(&thread->ts_choose_rs, 0, total_weight - 1); - do + /* + * If the latency limit is used, print a percentage of the current + * transaction latency from the latency limit. + */ + if (latency_limit) { - w -= sql_script[i++].weight; - } while (w >= 0); - - return i - 1; -} - -/* - * Allocate space for CState->prepared: we need one boolean for each command - * of each script. - */ -static void -allocCStatePrepared(CState *st) -{ - Assert(st->prepared == NULL); + pg_time_now_lazy(now); + appendPQExpBuffer(&buf, ", %.3f%% of the maximum time of tries was used", + (100.0 * (*now - st->txn_scheduled) / latency_limit)); + } + appendPQExpBufferStr(&buf, ")\n"); - st->prepared = pg_malloc_array(bool *, num_scripts); - for (int i = 0; i < num_scripts; i++) - { - ParsedScript *script = &sql_script[i]; - int numcmds; + pg_log_info("%s", buf.data); - for (numcmds = 0; script->commands[numcmds] != NULL; numcmds++) - ; - st->prepared[i] = pg_malloc0_array(bool, numcmds); - } + termPQExpBuffer(&buf); } /* - * Prepare the SQL command from st->use_file at command_num. + * Advance the state machine of a connection. */ static void -prepareCommand(CState *st, int command_num) +advanceConnectionState(TState *thread, CState *st, StatsData *agg) { - Command *command = sql_script[st->use_file].commands[command_num]; - - /* No prepare for non-SQL commands */ - if (command->type != SQL_COMMAND) - return; - if (!st->prepared) - allocCStatePrepared(st); + /* + * gettimeofday() isn't free, so we get the current timestamp lazily the + * first time it's needed, and reuse the same value throughout this + * function after that. This also ensures that e.g. the calculated + * latency reported in the log file and in the totals are the same. Zero + * means "not set yet". Reset "now" when we execute shell commands or + * expressions, which might take a non-negligible amount of time, though. + */ + pg_time_usec_t now = 0; - if (!st->prepared[st->use_file][command_num]) + /* + * Loop in the state machine, until we have to wait for a result from the + * server or have to sleep for throttling or \sleep. + * + * Note: In the switch-statement below, 'break' will loop back here, + * meaning "continue in the state machine". Return is used to return to + * the caller, giving the thread the opportunity to advance another + * client. + */ + for (;;) { - PGresult *res; + Command *command; - pg_log_debug("client %d preparing %s", st->id, command->prepname); - res = PQprepare(st->con, command->prepname, - command->argv[0], command->argc - 1, NULL); - if (PQresultStatus(res) != PGRES_COMMAND_OK) - pg_log_error("%s", PQerrorMessage(st->con)); - PQclear(res); - st->prepared[st->use_file][command_num] = true; - } -} + switch (st->state) + { + /* Select transaction (script) to run. */ + case CSTATE_CHOOSE_SCRIPT: + st->use_file = chooseScript(thread); + Assert(conditional_stack_empty(st->cstack)); -/* - * Prepare all the commands in the script that come after the \startpipeline - * that's at position st->command, and the first \endpipeline we find. - * - * This sets the ->prepared flag for each relevant command as well as the - * \startpipeline itself, but doesn't move the st->command counter. - */ -static void -prepareCommandsInPipeline(CState *st) -{ - int j; - Command **commands = sql_script[st->use_file].commands; + /* reset transaction variables to default values */ + st->estatus = ESTATUS_NO_ERROR; + st->tries = 1; - Assert(commands[st->command]->type == META_COMMAND && - commands[st->command]->meta == META_STARTPIPELINE); + pg_log_debug("client %d executing script \"%s\"", + st->id, sql_script[st->use_file].desc); - if (!st->prepared) - allocCStatePrepared(st); + /* + * If time is over, we're done; otherwise, get ready to start + * a new transaction, or to get throttled if that's requested. + */ + st->state = timer_exceeded ? CSTATE_FINISHED : + throttle_delay > 0 ? CSTATE_PREPARE_THROTTLE : CSTATE_START_TX; + break; - /* - * We set the 'prepared' flag on the \startpipeline itself to flag that we - * don't need to do this next time without calling prepareCommand(), even - * though we don't actually prepare this command. - */ - if (st->prepared[st->use_file][st->command]) - return; + /* Start new transaction (script) */ + case CSTATE_START_TX: + pg_time_now_lazy(&now); - for (j = st->command + 1; commands[j] != NULL; j++) - { - if (commands[j]->type == META_COMMAND && - commands[j]->meta == META_ENDPIPELINE) - break; + /* establish connection if needed, i.e. under --connect */ + if (st->con == NULL) + { + pg_time_usec_t start = now; - prepareCommand(st, j); - } + if ((st->con = doConnect()) == NULL) + { + /* + * as the bench is already running, we do not abort + * the process + */ + pg_log_error("client %d aborted while establishing connection", st->id); + st->state = CSTATE_ABORTED; + break; + } - st->prepared[st->use_file][st->command] = true; -} + /* reset now after connection */ + now = pg_time_now(); -/* Send a SQL command, using the chosen querymode */ -static bool -sendCommand(CState *st, Command *command) -{ - int r; + thread->conn_duration += now - start; - if (querymode == QUERY_SIMPLE) - { - char *sql; + /* Reset session-local state */ + pg_free(st->prepared); + st->prepared = NULL; + } - sql = pg_strdup(command->argv[0]); - sql = assignVariables(&st->variables, sql); + /* + * It is the first try to run this transaction. Remember the + * random state: maybe it will get an error and we will need + * to run it again. + */ + st->random_state = st->cs_func_rs; - pg_log_debug("client %d sending %s", st->id, sql); - r = PQsendQuery(st->con, sql); - pg_free(sql); - } - else if (querymode == QUERY_EXTENDED) - { - const char *sql = command->argv[0]; - const char *params[MAX_ARGS]; + /* record transaction start time */ + st->txn_begin = now; - getQueryParams(&st->variables, command, params); + /* + * When not throttling, this is also the transaction's + * scheduled start time. + */ + if (!throttle_delay) + st->txn_scheduled = now; - pg_log_debug("client %d sending %s", st->id, sql); - r = PQsendQueryParams(st->con, sql, command->argc - 1, - NULL, params, NULL, NULL, 0); - } - else if (querymode == QUERY_PREPARED) - { - const char *params[MAX_ARGS]; + /* Begin with the first command */ + st->state = CSTATE_START_COMMAND; + st->command = 0; + break; - prepareCommand(st, st->command); - getQueryParams(&st->variables, command, params); + /* + * Handle throttling once per transaction by sleeping. + */ + case CSTATE_PREPARE_THROTTLE: - pg_log_debug("client %d sending %s", st->id, command->prepname); - r = PQsendQueryPrepared(st->con, command->prepname, command->argc - 1, - params, NULL, NULL, 0); - } - else /* unknown sql mode */ - r = 0; + /* + * Generate a delay such that the series of delays will + * approximate a Poisson distribution centered on the + * throttle_delay time. + * + * If transactions are too slow or a given wait is shorter + * than a transaction, the next transaction will start right + * away. + */ + Assert(throttle_delay > 0); - if (r == 0) - { - pg_log_debug("client %d could not send %s", st->id, command->argv[0]); - return false; - } - else - return true; -} + thread->throttle_trigger += + getPoissonRand(&thread->ts_throttle_rs, throttle_delay); + st->txn_scheduled = thread->throttle_trigger; -/* - * Read and discard all available results from the connection. - */ -static void -discardAvailableResults(CState *st) -{ - PGresult *res = NULL; - - for (;;) - { - res = PQgetResult(st->con); - - /* - * Read and discard results until PQgetResult() returns NULL (no more - * results) or a connection failure is detected. If the pipeline - * status is PQ_PIPELINE_ABORTED, more results may still be available - * even after PQgetResult() returns NULL, so continue reading in that - * case. - */ - if ((res == NULL && PQpipelineStatus(st->con) != PQ_PIPELINE_ABORTED) || - PQstatus(st->con) == CONNECTION_BAD) - break; - - PQclear(res); - } - PQclear(res); -} + /* + * If --latency-limit is used, and this slot is already late + * so that the transaction will miss the latency limit even if + * it completed immediately, skip this time slot and loop to + * reschedule. + */ + if (latency_limit) + { + pg_time_now_lazy(&now); -/* - * Determine the error status based on the connection status and error code. - */ -static EStatus -getSQLErrorStatus(CState *st, const char *sqlState) -{ - discardAvailableResults(st); - if (PQstatus(st->con) == CONNECTION_BAD) - return ESTATUS_CONN_ERROR; + if (thread->throttle_trigger < now - latency_limit) + { + processXactStats(thread, st, &now, true, agg); - if (sqlState != NULL) - { - if (strcmp(sqlState, ERRCODE_T_R_SERIALIZATION_FAILURE) == 0) - return ESTATUS_SERIALIZATION_ERROR; - else if (strcmp(sqlState, ERRCODE_T_R_DEADLOCK_DETECTED) == 0) - return ESTATUS_DEADLOCK_ERROR; - } + /* + * Finish client if -T or -t was exceeded. + * + * Stop counting skipped transactions under -T as soon + * as the timer is exceeded. Because otherwise it can + * take a very long time to count all of them + * especially when quite a lot of them happen with + * unrealistically high rate setting in -R, which + * would prevent pgbench from ending immediately. + * Because of this behavior, note that there is no + * guarantee that all skipped transactions are counted + * under -T though there is under -t. This is OK in + * practice because it's very unlikely to happen with + * realistic setting. + */ + if (timer_exceeded || (nxacts > 0 && st->cnt >= nxacts)) + st->state = CSTATE_FINISHED; - return ESTATUS_OTHER_SQL_ERROR; -} + /* Go back to top of loop with CSTATE_PREPARE_THROTTLE */ + break; + } + } -/* - * Returns true if this type of error can be retried. - */ -static bool -canRetryError(EStatus estatus) -{ - return (estatus == ESTATUS_SERIALIZATION_ERROR || - estatus == ESTATUS_DEADLOCK_ERROR); -} + /* + * stop client if next transaction is beyond pgbench end of + * execution; otherwise, throttle it. + */ + st->state = end_time > 0 && st->txn_scheduled > end_time ? + CSTATE_FINISHED : CSTATE_THROTTLE; + break; -/* - * Returns true if --continue-on-error is specified and this error allows - * processing to continue. - */ -static bool -canContinueOnError(EStatus estatus) -{ - return (continue_on_error && - estatus == ESTATUS_OTHER_SQL_ERROR); -} + /* + * Wait until it's time to start next transaction. + */ + case CSTATE_THROTTLE: + pg_time_now_lazy(&now); -/* - * Process query response from the backend. - * - * If varprefix is not NULL, it's the variable name prefix where to store - * the results of the *last* command (META_GSET) or *all* commands - * (META_ASET). - * - * Returns true if everything is A-OK, false if any error occurs. - */ -static bool -readCommandResponse(CState *st, MetaCommand meta, char *varprefix) -{ - PGresult *res; - PGresult *next_res; - int qrynum = 0; + if (now < st->txn_scheduled) + return; /* still sleeping, nothing to do here */ - /* - * varprefix should be set only with \gset or \aset, and \endpipeline and - * SQL commands do not need it. - */ - Assert((meta == META_NONE && varprefix == NULL) || - ((meta == META_ENDPIPELINE) && varprefix == NULL) || - ((meta == META_GSET || meta == META_ASET) && varprefix != NULL)); + /* done sleeping, but don't start transaction if we're done */ + st->state = timer_exceeded ? CSTATE_FINISHED : CSTATE_START_TX; + break; - res = PQgetResult(st->con); + /* + * Send a command to server (or execute a meta-command) + */ + case CSTATE_START_COMMAND: + command = sql_script[st->use_file].commands[st->command]; - while (res != NULL) - { - bool is_last; + /* + * Transition to script end processing if done, but close up + * shop if a pipeline is open at this point. + */ + if (command == NULL) + { + if (PQpipelineStatus(st->con) == PQ_PIPELINE_OFF) + st->state = CSTATE_END_TX; + else + { + pg_log_error("client %d aborted: end of script reached with pipeline open", + st->id); + st->state = CSTATE_ABORTED; + } - /* peek at the next result to know whether the current is last */ - next_res = PQgetResult(st->con); - is_last = (next_res == NULL); + break; + } - switch (PQresultStatus(res)) - { - case PGRES_COMMAND_OK: /* non-SELECT commands */ - case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */ - if (is_last && meta == META_GSET) + /* record begin time of next command, and initiate it */ + if (report_per_command) { - pg_log_error("client %d script %d command %d query %d: expected one row, got %d", - st->id, st->use_file, st->command, qrynum, 0); - st->estatus = ESTATUS_META_COMMAND_ERROR; - goto error; + pg_time_now_lazy(&now); + st->stmt_begin = now; } - break; - case PGRES_TUPLES_OK: - if ((is_last && meta == META_GSET) || meta == META_ASET) + /* Execute the command */ + if (command->type == SQL_COMMAND) { - int ntuples = PQntuples(res); - const char *context = (meta == META_ASET) ? "aset" : "gset"; - - if (meta == META_GSET && ntuples != 1) - { - /* under \gset, report the error */ - pg_log_error("client %d script %d command %d query %d: expected one row, got %d", - st->id, st->use_file, st->command, qrynum, PQntuples(res)); - st->estatus = ESTATUS_META_COMMAND_ERROR; - goto error; - } - else if (meta == META_ASET && ntuples <= 0) - { - /* coldly skip empty result under \aset */ - break; - } - - /* store results into variables */ - for (int fld = 0; fld < PQnfields(res); fld++) + /* disallow \aset and \gset in pipeline mode */ + if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) { - char *varname = PQfname(res, fld); - bool ok; - - /* allocate varname only if necessary, freed below */ - if (*varprefix != '\0') - varname = psprintf("%s%s", varprefix, varname); - - /* store last row result */ - if (PQgetisnull(res, ntuples - 1, fld)) + if (command->meta == META_GSET) { - PgBenchValue nullval; - - setNullValue(&nullval); - ok = putVariableValue(&st->variables, context, - varname, &nullval); + commandFailed(st, "gset", "\\gset is not allowed in pipeline mode"); + st->state = CSTATE_ABORTED; + break; } - else - ok = putVariable(&st->variables, context, varname, - PQgetvalue(res, ntuples - 1, fld)); - - if (!ok) + else if (command->meta == META_ASET) { - /* internal error */ - pg_log_error("client %d script %d command %d query %d: error storing into variable %s", - st->id, st->use_file, st->command, qrynum, varname); - st->estatus = ESTATUS_META_COMMAND_ERROR; - goto error; + commandFailed(st, "aset", "\\aset is not allowed in pipeline mode"); + st->state = CSTATE_ABORTED; + break; } + } - if (*varprefix != '\0') - pfree(varname); + if (!sendCommand(st, command)) + { + commandFailed(st, "SQL", "SQL command send failed"); + st->state = CSTATE_ABORTED; + } + else + { + /* Wait for results, unless in pipeline mode */ + if (PQpipelineStatus(st->con) == PQ_PIPELINE_OFF) + st->state = CSTATE_WAIT_RESULT; + else + st->state = CSTATE_END_COMMAND; } } - /* otherwise the result is simply thrown away by PQclear below */ - break; - - case PGRES_PIPELINE_SYNC: - pg_log_debug("client %d pipeline ending, ongoing syncs: %d", - st->id, st->num_syncs); - st->num_syncs--; - if (st->num_syncs == 0 && PQexitPipelineMode(st->con) != 1) - pg_log_error("client %d failed to exit pipeline mode: %s", st->id, - PQresultErrorMessage(res)); - break; - - case PGRES_COPY_IN: - case PGRES_COPY_OUT: - case PGRES_COPY_BOTH: - pg_log_error("COPY is not supported in pgbench, aborting"); + else if (command->type == META_COMMAND) + { + /*----- + * Possible state changes when executing meta commands: + * - on errors CSTATE_ABORTED + * - on sleep CSTATE_SLEEP + * - else CSTATE_END_COMMAND + */ + st->state = executeMetaCommand(st, &now); + if (st->state == CSTATE_ABORTED) + st->estatus = ESTATUS_META_COMMAND_ERROR; + } /* - * We need to exit the copy state. Otherwise, PQgetResult() - * will always return an empty PGresult as an effect of - * getCopyResult(), leading to an infinite loop in the error - * cleanup done below. + * We're now waiting for an SQL command to complete, or + * finished processing a metacommand, or need to sleep, or + * something bad happened. */ - PQendcopy(st->con); - goto error; + Assert(st->state == CSTATE_WAIT_RESULT || + st->state == CSTATE_END_COMMAND || + st->state == CSTATE_SLEEP || + st->state == CSTATE_ABORTED); + break; - case PGRES_NONFATAL_ERROR: - case PGRES_FATAL_ERROR: - st->estatus = getSQLErrorStatus(st, PQresultErrorField(res, - PG_DIAG_SQLSTATE)); - if (canRetryError(st->estatus) || canContinueOnError(st->estatus)) + /* + * non executed conditional branch + */ + case CSTATE_SKIP_COMMAND: + Assert(!conditional_active(st->cstack)); + /* quickly skip commands until something to do... */ + while (true) { - if (verbose_errors) - commandError(st, PQresultErrorMessage(res)); - goto error; - } - pg_fallthrough; - - default: - /* anything else is unexpected */ - pg_log_error("client %d script %d aborted in command %d query %d: %s", - st->id, st->use_file, st->command, qrynum, - PQresultErrorMessage(res)); - goto error; - } - - PQclear(res); - qrynum++; - res = next_res; - } - - if (qrynum == 0) - { - pg_log_error("client %d command %d: no results", st->id, st->command); - return false; - } - - return true; - -error: - PQclear(res); - PQclear(next_res); - discardAvailableResults(st); - - return false; -} - -/* - * Parse the argument to a \sleep command, and return the requested amount - * of delay, in microseconds. Returns true on success, false on error. - */ -static bool -evaluateSleep(Variables *variables, int argc, char **argv, int *usecs) -{ - char *var; - int usec; - - if (*argv[1] == ':') - { - if ((var = getVariable(variables, argv[1] + 1)) == NULL) - { - 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]); - - if (argc > 2) - { - if (pg_strcasecmp(argv[2], "ms") == 0) - usec *= 1000; - else if (pg_strcasecmp(argv[2], "s") == 0) - usec *= 1000000; - } - else - usec *= 1000000; - - *usecs = usec; - return true; -} - - -/* - * Returns true if the error can be retried. - */ -static bool -doRetry(CState *st, pg_time_usec_t *now) -{ - Assert(st->estatus != ESTATUS_NO_ERROR); - - /* We can only retry serialization or deadlock errors. */ - if (!canRetryError(st->estatus)) - return false; - - /* - * We must have at least one option to limit the retrying of transactions - * that got an error. - */ - Assert(max_tries || latency_limit || duration > 0); - - /* - * We cannot retry the error if we have reached the maximum number of - * tries. - */ - if (max_tries && st->tries >= max_tries) - return false; - - /* - * We cannot retry the error if we spent too much time on this - * transaction. - */ - if (latency_limit) - { - pg_time_now_lazy(now); - if (*now - st->txn_scheduled > latency_limit) - return false; - } - - /* - * We cannot retry the error if the benchmark duration is over. - */ - if (timer_exceeded) - return false; - - /* OK */ - return true; -} - -/* - * Read and discard results until the last sync point. - */ -static int -discardUntilSync(CState *st) -{ - bool received_sync = false; - - /* - * Send a Sync message to ensure at least one PGRES_PIPELINE_SYNC is - * received and to avoid an infinite loop, since all earlier ones may have - * already been received. - */ - if (!PQpipelineSync(st->con)) - { - pg_log_error("client %d aborted: failed to send a pipeline sync", - st->id); - return 0; - } + command = sql_script[st->use_file].commands[st->command]; - /* - * Continue reading results until the last sync point, i.e., until - * reaching null just after PGRES_PIPELINE_SYNC. - */ - for (;;) - { - PGresult *res = PQgetResult(st->con); + /* cannot reach end of script in that state */ + Assert(command != NULL); - if (PQstatus(st->con) == CONNECTION_BAD) - { - pg_log_error("client %d aborted while rolling back the transaction after an error; perhaps the backend died while processing", - st->id); - PQclear(res); - return 0; - } + /* + * if this is conditional related, update conditional + * state + */ + if (command->type == META_COMMAND && + (command->meta == META_IF || + command->meta == META_ELIF || + command->meta == META_ELSE || + command->meta == META_ENDIF)) + { + switch (conditional_stack_peek(st->cstack)) + { + case IFSTATE_FALSE: + if (command->meta == META_IF) + { + /* nested if in skipped branch - ignore */ + conditional_stack_push(st->cstack, + IFSTATE_IGNORED); + st->command++; + } + else if (command->meta == META_ELIF) + { + /* we must evaluate the condition */ + st->state = CSTATE_START_COMMAND; + } + else if (command->meta == META_ELSE) + { + /* we must execute next command */ + conditional_stack_poke(st->cstack, + IFSTATE_ELSE_TRUE); + st->state = CSTATE_START_COMMAND; + st->command++; + } + else if (command->meta == META_ENDIF) + { + Assert(!conditional_stack_empty(st->cstack)); + conditional_stack_pop(st->cstack); + if (conditional_active(st->cstack)) + st->state = CSTATE_START_COMMAND; + /* else state remains CSTATE_SKIP_COMMAND */ + st->command++; + } + break; - if (PQresultStatus(res) == PGRES_PIPELINE_SYNC) - received_sync = true; - else if (received_sync && res == NULL) - { - /* - * Reset ongoing sync count to 0 since all PGRES_PIPELINE_SYNC - * results have been discarded. - */ - st->num_syncs = 0; - break; - } - else - { - /* - * If a PGRES_PIPELINE_SYNC is followed by something other than - * PGRES_PIPELINE_SYNC or NULL, another PGRES_PIPELINE_SYNC will - * appear later. Reset received_sync to false to wait for it. - */ - received_sync = false; - } - PQclear(res); - } - - /* exit pipeline */ - if (PQexitPipelineMode(st->con) != 1) - { - pg_log_error("client %d aborted: failed to exit pipeline mode for rolling back the failed transaction", - st->id); - return 0; - } - return 1; -} - -/* - * Get the transaction status at the end of a command especially for - * checking if we are in a (failed) transaction block. - */ -static TStatus -getTransactionStatus(PGconn *con) -{ - PGTransactionStatusType tx_status; - - tx_status = PQtransactionStatus(con); - switch (tx_status) - { - case PQTRANS_IDLE: - return TSTATUS_IDLE; - case PQTRANS_INTRANS: - case PQTRANS_INERROR: - return TSTATUS_IN_BLOCK; - case PQTRANS_UNKNOWN: - /* PQTRANS_UNKNOWN is expected given a broken connection */ - if (PQstatus(con) == CONNECTION_BAD) - return TSTATUS_CONN_ERROR; - pg_fallthrough; - case PQTRANS_ACTIVE: - default: - - /* - * We cannot find out whether we are in a transaction block or - * not. Internal error which should never occur. - */ - pg_log_error("unexpected transaction status %d", tx_status); - return TSTATUS_OTHER_ERROR; - } - - /* not reached */ - Assert(false); - return TSTATUS_OTHER_ERROR; -} - -/* - * Print verbose messages of an error - */ -static void -printVerboseErrorMessages(CState *st, pg_time_usec_t *now, bool is_retry) -{ - PQExpBufferData buf; - - initPQExpBuffer(&buf); - - printfPQExpBuffer(&buf, "client %d ", st->id); - appendPQExpBufferStr(&buf, (is_retry ? - "repeats the transaction after the error" : - "ends the failed transaction")); - appendPQExpBuffer(&buf, " (try %u", st->tries); - - /* Print max_tries if it is not unlimited. */ - if (max_tries) - appendPQExpBuffer(&buf, "/%u", max_tries); - - /* - * If the latency limit is used, print a percentage of the current - * transaction latency from the latency limit. - */ - if (latency_limit) - { - pg_time_now_lazy(now); - appendPQExpBuffer(&buf, ", %.3f%% of the maximum time of tries was used", - (100.0 * (*now - st->txn_scheduled) / latency_limit)); - } - appendPQExpBufferStr(&buf, ")\n"); - - pg_log_info("%s", buf.data); - - termPQExpBuffer(&buf); -} - -/* - * Advance the state machine of a connection. - */ -static void -advanceConnectionState(TState *thread, CState *st, StatsData *agg) -{ - - /* - * gettimeofday() isn't free, so we get the current timestamp lazily the - * first time it's needed, and reuse the same value throughout this - * function after that. This also ensures that e.g. the calculated - * latency reported in the log file and in the totals are the same. Zero - * means "not set yet". Reset "now" when we execute shell commands or - * expressions, which might take a non-negligible amount of time, though. - */ - pg_time_usec_t now = 0; - - /* - * Loop in the state machine, until we have to wait for a result from the - * server or have to sleep for throttling or \sleep. - * - * Note: In the switch-statement below, 'break' will loop back here, - * meaning "continue in the state machine". Return is used to return to - * the caller, giving the thread the opportunity to advance another - * client. - */ - for (;;) - { - Command *command; - - switch (st->state) - { - /* Select transaction (script) to run. */ - case CSTATE_CHOOSE_SCRIPT: - st->use_file = chooseScript(thread); - Assert(conditional_stack_empty(st->cstack)); - - /* reset transaction variables to default values */ - st->estatus = ESTATUS_NO_ERROR; - st->tries = 1; - - pg_log_debug("client %d executing script \"%s\"", - st->id, sql_script[st->use_file].desc); - - /* - * If time is over, we're done; otherwise, get ready to start - * a new transaction, or to get throttled if that's requested. - */ - st->state = timer_exceeded ? CSTATE_FINISHED : - throttle_delay > 0 ? CSTATE_PREPARE_THROTTLE : CSTATE_START_TX; - break; - - /* Start new transaction (script) */ - case CSTATE_START_TX: - pg_time_now_lazy(&now); - - /* establish connection if needed, i.e. under --connect */ - if (st->con == NULL) - { - pg_time_usec_t start = now; - - if ((st->con = doConnect()) == NULL) - { - /* - * as the bench is already running, we do not abort - * the process - */ - pg_log_error("client %d aborted while establishing connection", st->id); - st->state = CSTATE_ABORTED; - break; - } - - /* reset now after connection */ - now = pg_time_now(); - - thread->conn_duration += now - start; - - /* Reset session-local state */ - pg_free(st->prepared); - st->prepared = NULL; - } - - /* - * It is the first try to run this transaction. Remember the - * random state: maybe it will get an error and we will need - * to run it again. - */ - st->random_state = st->cs_func_rs; - - /* record transaction start time */ - st->txn_begin = now; - - /* - * When not throttling, this is also the transaction's - * scheduled start time. - */ - if (!throttle_delay) - st->txn_scheduled = now; - - /* Begin with the first command */ - st->state = CSTATE_START_COMMAND; - st->command = 0; - break; - - /* - * Handle throttling once per transaction by sleeping. - */ - case CSTATE_PREPARE_THROTTLE: - - /* - * Generate a delay such that the series of delays will - * approximate a Poisson distribution centered on the - * throttle_delay time. - * - * If transactions are too slow or a given wait is shorter - * than a transaction, the next transaction will start right - * away. - */ - Assert(throttle_delay > 0); - - thread->throttle_trigger += - getPoissonRand(&thread->ts_throttle_rs, throttle_delay); - st->txn_scheduled = thread->throttle_trigger; - - /* - * If --latency-limit is used, and this slot is already late - * so that the transaction will miss the latency limit even if - * it completed immediately, skip this time slot and loop to - * reschedule. - */ - if (latency_limit) - { - pg_time_now_lazy(&now); - - if (thread->throttle_trigger < now - latency_limit) - { - processXactStats(thread, st, &now, true, agg); - - /* - * Finish client if -T or -t was exceeded. - * - * Stop counting skipped transactions under -T as soon - * as the timer is exceeded. Because otherwise it can - * take a very long time to count all of them - * especially when quite a lot of them happen with - * unrealistically high rate setting in -R, which - * would prevent pgbench from ending immediately. - * Because of this behavior, note that there is no - * guarantee that all skipped transactions are counted - * under -T though there is under -t. This is OK in - * practice because it's very unlikely to happen with - * realistic setting. - */ - if (timer_exceeded || (nxacts > 0 && st->cnt >= nxacts)) - st->state = CSTATE_FINISHED; - - /* Go back to top of loop with CSTATE_PREPARE_THROTTLE */ - break; - } - } - - /* - * stop client if next transaction is beyond pgbench end of - * execution; otherwise, throttle it. - */ - st->state = end_time > 0 && st->txn_scheduled > end_time ? - CSTATE_FINISHED : CSTATE_THROTTLE; - break; - - /* - * Wait until it's time to start next transaction. - */ - case CSTATE_THROTTLE: - pg_time_now_lazy(&now); - - if (now < st->txn_scheduled) - return; /* still sleeping, nothing to do here */ - - /* done sleeping, but don't start transaction if we're done */ - st->state = timer_exceeded ? CSTATE_FINISHED : CSTATE_START_TX; - break; - - /* - * Send a command to server (or execute a meta-command) - */ - case CSTATE_START_COMMAND: - command = sql_script[st->use_file].commands[st->command]; - - /* - * Transition to script end processing if done, but close up - * shop if a pipeline is open at this point. - */ - if (command == NULL) - { - if (PQpipelineStatus(st->con) == PQ_PIPELINE_OFF) - st->state = CSTATE_END_TX; - else - { - pg_log_error("client %d aborted: end of script reached with pipeline open", - st->id); - st->state = CSTATE_ABORTED; - } - - break; - } - - /* record begin time of next command, and initiate it */ - if (report_per_command) - { - pg_time_now_lazy(&now); - st->stmt_begin = now; - } - - /* Execute the command */ - if (command->type == SQL_COMMAND) - { - /* disallow \aset and \gset in pipeline mode */ - if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) - { - if (command->meta == META_GSET) - { - commandFailed(st, "gset", "\\gset is not allowed in pipeline mode"); - st->state = CSTATE_ABORTED; - break; - } - else if (command->meta == META_ASET) - { - commandFailed(st, "aset", "\\aset is not allowed in pipeline mode"); - st->state = CSTATE_ABORTED; - break; - } - } - - if (!sendCommand(st, command)) - { - commandFailed(st, "SQL", "SQL command send failed"); - st->state = CSTATE_ABORTED; - } - else - { - /* Wait for results, unless in pipeline mode */ - if (PQpipelineStatus(st->con) == PQ_PIPELINE_OFF) - st->state = CSTATE_WAIT_RESULT; - else - st->state = CSTATE_END_COMMAND; - } - } - else if (command->type == META_COMMAND) - { - /*----- - * Possible state changes when executing meta commands: - * - on errors CSTATE_ABORTED - * - on sleep CSTATE_SLEEP - * - else CSTATE_END_COMMAND - */ - st->state = executeMetaCommand(st, &now); - if (st->state == CSTATE_ABORTED) - st->estatus = ESTATUS_META_COMMAND_ERROR; - } - - /* - * We're now waiting for an SQL command to complete, or - * finished processing a metacommand, or need to sleep, or - * something bad happened. - */ - Assert(st->state == CSTATE_WAIT_RESULT || - st->state == CSTATE_END_COMMAND || - st->state == CSTATE_SLEEP || - st->state == CSTATE_ABORTED); - break; - - /* - * non executed conditional branch - */ - case CSTATE_SKIP_COMMAND: - Assert(!conditional_active(st->cstack)); - /* quickly skip commands until something to do... */ - while (true) - { - command = sql_script[st->use_file].commands[st->command]; - - /* cannot reach end of script in that state */ - Assert(command != NULL); - - /* - * if this is conditional related, update conditional - * state - */ - if (command->type == META_COMMAND && - (command->meta == META_IF || - command->meta == META_ELIF || - command->meta == META_ELSE || - command->meta == META_ENDIF)) - { - switch (conditional_stack_peek(st->cstack)) - { - case IFSTATE_FALSE: - if (command->meta == META_IF) - { - /* nested if in skipped branch - ignore */ - conditional_stack_push(st->cstack, - IFSTATE_IGNORED); - st->command++; - } - else if (command->meta == META_ELIF) - { - /* we must evaluate the condition */ - st->state = CSTATE_START_COMMAND; - } - else if (command->meta == META_ELSE) - { - /* we must execute next command */ - conditional_stack_poke(st->cstack, - IFSTATE_ELSE_TRUE); - st->state = CSTATE_START_COMMAND; - st->command++; - } - else if (command->meta == META_ENDIF) - { - Assert(!conditional_stack_empty(st->cstack)); - conditional_stack_pop(st->cstack); - if (conditional_active(st->cstack)) - st->state = CSTATE_START_COMMAND; - /* else state remains CSTATE_SKIP_COMMAND */ - st->command++; - } - break; - - case IFSTATE_IGNORED: - case IFSTATE_ELSE_FALSE: - if (command->meta == META_IF) - conditional_stack_push(st->cstack, - IFSTATE_IGNORED); - else if (command->meta == META_ENDIF) - { - Assert(!conditional_stack_empty(st->cstack)); - conditional_stack_pop(st->cstack); - if (conditional_active(st->cstack)) - st->state = CSTATE_START_COMMAND; - } - /* could detect "else" & "elif" after "else" */ - st->command++; - break; - - case IFSTATE_NONE: - case IFSTATE_TRUE: - case IFSTATE_ELSE_TRUE: - default: - - /* - * inconsistent if inactive, unreachable dead - * code - */ - Assert(false); - } - } - else - { - /* skip and consider next */ - st->command++; - } - - if (st->state != CSTATE_SKIP_COMMAND) - /* out of quick skip command loop */ - break; - } - break; - - /* - * Wait for the current SQL command to complete - */ - case CSTATE_WAIT_RESULT: - pg_log_debug("client %d receiving", st->id); - - /* - * Only check for new network data if we processed all data - * fetched prior. Otherwise we end up doing a syscall for each - * individual pipelined query, which has a measurable - * performance impact. - */ - if (PQisBusy(st->con) && !PQconsumeInput(st->con)) - { - /* there's something wrong */ - commandFailed(st, "SQL", "perhaps the backend died while processing"); - st->state = CSTATE_ABORTED; - break; - } - if (PQisBusy(st->con)) - return; /* don't have the whole result yet */ - - /* store or discard the query results */ - if (readCommandResponse(st, - sql_script[st->use_file].commands[st->command]->meta, - sql_script[st->use_file].commands[st->command]->varprefix)) - { - /* - * outside of pipeline mode: stop reading results. - * pipeline mode: continue reading results until an - * end-of-pipeline response. - */ - if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) - st->state = CSTATE_END_COMMAND; - } - else if (canRetryError(st->estatus) || canContinueOnError(st->estatus)) - st->state = CSTATE_ERROR; - else - st->state = CSTATE_ABORTED; - break; - - /* - * Wait until sleep is done. This state is entered after a - * \sleep metacommand. The behavior is similar to - * CSTATE_THROTTLE, but proceeds to CSTATE_START_COMMAND - * instead of CSTATE_START_TX. - */ - case CSTATE_SLEEP: - pg_time_now_lazy(&now); - if (now < st->sleep_until) - return; /* still sleeping, nothing to do here */ - /* Else done sleeping. */ - st->state = CSTATE_END_COMMAND; - break; - - /* - * End of command: record stats and proceed to next command. - */ - case CSTATE_END_COMMAND: - - /* - * command completed: accumulate per-command execution times - * in thread-local data structure, if per-command latencies - * are requested. - */ - if (report_per_command) - { - pg_time_now_lazy(&now); - - command = sql_script[st->use_file].commands[st->command]; - /* XXX could use a mutex here, but we choose not to */ - addToSimpleStats(&command->stats, - PG_TIME_GET_DOUBLE(now - st->stmt_begin)); - } - - /* Go ahead with next command, to be executed or skipped */ - st->command++; - st->state = conditional_active(st->cstack) ? - CSTATE_START_COMMAND : CSTATE_SKIP_COMMAND; - break; - - /* - * Clean up after an error. - */ - case CSTATE_ERROR: - { - TStatus tstatus; - - Assert(st->estatus != ESTATUS_NO_ERROR); - - /* Clear the conditional stack */ - conditional_stack_reset(st->cstack); - - /* Read and discard until a sync point in pipeline mode */ - if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) - { - if (!discardUntilSync(st)) - { - st->state = CSTATE_ABORTED; - break; - } - } - - /* - * Check if we have a (failed) transaction block or not, - * and roll it back if any. - */ - tstatus = getTransactionStatus(st->con); - if (tstatus == TSTATUS_IN_BLOCK) - { - /* Try to rollback a (failed) transaction block. */ - if (!PQsendQuery(st->con, "ROLLBACK")) - { - pg_log_error("client %d aborted: failed to send sql command for rolling back the failed transaction", - st->id); - st->state = CSTATE_ABORTED; - } - else - st->state = CSTATE_WAIT_ROLLBACK_RESULT; - } - else if (tstatus == TSTATUS_IDLE) - { - /* - * If time is over, we're done; otherwise, check if we - * can retry the error. - */ - st->state = timer_exceeded ? CSTATE_FINISHED : - doRetry(st, &now) ? CSTATE_RETRY : CSTATE_FAILURE; - } - else - { - if (tstatus == TSTATUS_CONN_ERROR) - pg_log_error("perhaps the backend died while processing"); - - pg_log_error("client %d aborted while receiving the transaction status", st->id); - st->state = CSTATE_ABORTED; - } - break; - } - - /* - * Wait for the rollback command to complete - */ - case CSTATE_WAIT_ROLLBACK_RESULT: - { - PGresult *res; - - pg_log_debug("client %d receiving", st->id); - if (!PQconsumeInput(st->con)) - { - pg_log_error("client %d aborted while rolling back the transaction after an error; perhaps the backend died while processing", - st->id); - st->state = CSTATE_ABORTED; - break; - } - if (PQisBusy(st->con)) - return; /* don't have the whole result yet */ - - /* - * Read and discard the query result; - */ - res = PQgetResult(st->con); - switch (PQresultStatus(res)) - { - case PGRES_COMMAND_OK: - /* OK */ - PQclear(res); - /* null must be returned */ - res = PQgetResult(st->con); - Assert(res == NULL); - - /* - * If time is over, we're done; otherwise, check - * if we can retry the error. - */ - st->state = timer_exceeded ? CSTATE_FINISHED : - doRetry(st, &now) ? CSTATE_RETRY : CSTATE_FAILURE; - break; - default: - pg_log_error("client %d aborted while rolling back the transaction after an error; %s", - st->id, PQerrorMessage(st->con)); - PQclear(res); - st->state = CSTATE_ABORTED; - break; - } - break; - } - - /* - * Retry the transaction after an error. - */ - case CSTATE_RETRY: - command = sql_script[st->use_file].commands[st->command]; - - /* - * Inform that the transaction will be retried after the - * error. - */ - if (verbose_errors) - printVerboseErrorMessages(st, &now, true); - - /* Count tries and retries */ - st->tries++; - command->retries++; - - /* - * Reset the random state as they were at the beginning of the - * transaction. - */ - st->cs_func_rs = st->random_state; - - /* Process the first transaction command. */ - st->command = 0; - st->estatus = ESTATUS_NO_ERROR; - st->state = CSTATE_START_COMMAND; - break; - - /* - * Record a failed transaction. - */ - case CSTATE_FAILURE: - command = sql_script[st->use_file].commands[st->command]; - - /* Accumulate the failure. */ - command->failures++; - - /* - * Inform that the failed transaction will not be retried. - */ - if (verbose_errors) - printVerboseErrorMessages(st, &now, false); - - /* End the failed transaction. */ - st->state = CSTATE_END_TX; - break; - - /* - * End of transaction (end of script, really). - */ - case CSTATE_END_TX: - { - TStatus tstatus; - - /* transaction finished: calculate latency and do log */ - processXactStats(thread, st, &now, false, agg); - - /* - * missing \endif... cannot happen if CheckConditional was - * okay - */ - Assert(conditional_stack_empty(st->cstack)); - - /* - * We must complete all the transaction blocks that were - * started in this script. - */ - tstatus = getTransactionStatus(st->con); - if (tstatus == TSTATUS_IN_BLOCK) - { - pg_log_error("client %d aborted: end of script reached without completing the last transaction", - st->id); - st->state = CSTATE_ABORTED; - break; - } - else if (tstatus != TSTATUS_IDLE) - { - if (tstatus == TSTATUS_CONN_ERROR) - pg_log_error("perhaps the backend died while processing"); - - pg_log_error("client %d aborted while receiving the transaction status", st->id); - st->state = CSTATE_ABORTED; - break; - } - - if (is_connect) - { - pg_time_usec_t start = now; - - pg_time_now_lazy(&start); - finishCon(st); - now = pg_time_now(); - thread->conn_duration += now - start; - } - - if ((st->cnt >= nxacts && duration <= 0) || timer_exceeded) - { - /* script completed */ - st->state = CSTATE_FINISHED; - break; - } - - /* next transaction (script) */ - st->state = CSTATE_CHOOSE_SCRIPT; - - /* - * Ensure that we always return on this point, so as to - * avoid an infinite loop if the script only contains meta - * commands. - */ - return; - } - - /* - * Final states. Close the connection if it's still open. - */ - case CSTATE_ABORTED: - case CSTATE_FINISHED: - - /* - * Don't measure the disconnection delays here even if in - * CSTATE_FINISHED and -C/--connect option is specified. - * Because in this case all the connections that this thread - * established are closed at the end of transactions and the - * disconnection delays should have already been measured at - * that moment. - * - * In CSTATE_ABORTED state, the measurement is no longer - * necessary because we cannot report complete results anyways - * in this case. - */ - finishCon(st); - return; - } - } -} - -/* - * Subroutine for advanceConnectionState -- initiate or execute the current - * meta command, and return the next state to set. - * - * *now is updated to the current time, unless the command is expected to - * take no time to execute. - */ -static ConnectionStateEnum -executeMetaCommand(CState *st, pg_time_usec_t *now) -{ - Command *command = sql_script[st->use_file].commands[st->command]; - int argc; - char **argv; - - Assert(command != NULL && command->type == META_COMMAND); - - argc = command->argc; - argv = command->argv; - - if (unlikely(__pg_log_level <= PG_LOG_DEBUG)) - { - PQExpBufferData buf; - - initPQExpBuffer(&buf); - - printfPQExpBuffer(&buf, "client %d executing \\%s", st->id, argv[0]); - for (int i = 1; i < argc; i++) - appendPQExpBuffer(&buf, " %s", argv[i]); - - pg_log_debug("%s", buf.data); - - termPQExpBuffer(&buf); - } - - if (command->meta == META_SLEEP) - { - int usec; - - /* - * A \sleep doesn't execute anything, we just get the delay from the - * argument, and enter the CSTATE_SLEEP state. (The per-command - * latency will be recorded in CSTATE_SLEEP state, not here, after the - * delay has elapsed.) - */ - if (!evaluateSleep(&st->variables, argc, argv, &usec)) - { - commandFailed(st, "sleep", "execution of meta-command failed"); - return CSTATE_ABORTED; - } - - pg_time_now_lazy(now); - st->sleep_until = (*now) + usec; - return CSTATE_SLEEP; - } - else if (command->meta == META_SET) - { - PgBenchExpr *expr = command->expr; - PgBenchValue result; - - if (!evaluateExpr(st, expr, &result)) - { - commandFailed(st, argv[0], "evaluation of meta-command failed"); - return CSTATE_ABORTED; - } - - if (!putVariableValue(&st->variables, argv[0], argv[1], &result)) - { - commandFailed(st, "set", "assignment of meta-command failed"); - return CSTATE_ABORTED; - } - } - else if (command->meta == META_IF) - { - /* backslash commands with an expression to evaluate */ - PgBenchExpr *expr = command->expr; - PgBenchValue result; - bool cond; - - if (!evaluateExpr(st, expr, &result)) - { - commandFailed(st, argv[0], "evaluation of meta-command failed"); - return CSTATE_ABORTED; - } - - cond = valueTruth(&result); - conditional_stack_push(st->cstack, cond ? IFSTATE_TRUE : IFSTATE_FALSE); - } - else if (command->meta == META_ELIF) - { - /* backslash commands with an expression to evaluate */ - PgBenchExpr *expr = command->expr; - PgBenchValue result; - bool cond; - - if (conditional_stack_peek(st->cstack) == IFSTATE_TRUE) - { - /* elif after executed block, skip eval and wait for endif. */ - conditional_stack_poke(st->cstack, IFSTATE_IGNORED); - return CSTATE_END_COMMAND; - } - - if (!evaluateExpr(st, expr, &result)) - { - commandFailed(st, argv[0], "evaluation of meta-command failed"); - return CSTATE_ABORTED; - } - - cond = valueTruth(&result); - Assert(conditional_stack_peek(st->cstack) == IFSTATE_FALSE); - conditional_stack_poke(st->cstack, cond ? IFSTATE_TRUE : IFSTATE_FALSE); - } - else if (command->meta == META_ELSE) - { - switch (conditional_stack_peek(st->cstack)) - { - case IFSTATE_TRUE: - conditional_stack_poke(st->cstack, IFSTATE_ELSE_FALSE); - break; - case IFSTATE_FALSE: /* inconsistent if active */ - case IFSTATE_IGNORED: /* inconsistent if active */ - case IFSTATE_NONE: /* else without if */ - case IFSTATE_ELSE_TRUE: /* else after else */ - case IFSTATE_ELSE_FALSE: /* else after else */ - default: - /* dead code if conditional check is ok */ - Assert(false); - } - } - else if (command->meta == META_ENDIF) - { - Assert(!conditional_stack_empty(st->cstack)); - conditional_stack_pop(st->cstack); - } - else if (command->meta == META_SETSHELL) - { - if (!runShellCommand(&st->variables, argv[1], argv + 2, argc - 2)) - { - commandFailed(st, "setshell", "execution of meta-command failed"); - return CSTATE_ABORTED; - } - } - else if (command->meta == META_SHELL) - { - if (!runShellCommand(&st->variables, NULL, argv + 1, argc - 1)) - { - commandFailed(st, "shell", "execution of meta-command failed"); - return CSTATE_ABORTED; - } - } - else if (command->meta == META_STARTPIPELINE) - { - /* - * In pipeline mode, we use a workflow based on libpq pipeline - * functions. - */ - if (querymode == QUERY_SIMPLE) - { - commandFailed(st, "startpipeline", "cannot use pipeline mode with the simple query protocol"); - return CSTATE_ABORTED; - } - - /* - * If we're in prepared-query mode, we need to prepare all the - * commands that are inside the pipeline before we actually start the - * pipeline itself. This solves the problem that running BEGIN - * ISOLATION LEVEL SERIALIZABLE in a pipeline would fail due to a - * snapshot having been acquired by the prepare within the pipeline. - */ - if (querymode == QUERY_PREPARED) - prepareCommandsInPipeline(st); - - if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) - { - commandFailed(st, "startpipeline", "already in pipeline mode"); - return CSTATE_ABORTED; - } - if (PQenterPipelineMode(st->con) == 0) - { - commandFailed(st, "startpipeline", "failed to enter pipeline mode"); - return CSTATE_ABORTED; - } - } - else if (command->meta == META_SYNCPIPELINE) - { - if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) - { - commandFailed(st, "syncpipeline", "not in pipeline mode"); - return CSTATE_ABORTED; - } - if (PQsendPipelineSync(st->con) == 0) - { - commandFailed(st, "syncpipeline", "failed to send a pipeline sync"); - return CSTATE_ABORTED; - } - st->num_syncs++; - } - else if (command->meta == META_ENDPIPELINE) - { - if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) - { - commandFailed(st, "endpipeline", "not in pipeline mode"); - return CSTATE_ABORTED; - } - if (!PQpipelineSync(st->con)) - { - commandFailed(st, "endpipeline", "failed to send a pipeline sync"); - return CSTATE_ABORTED; - } - st->num_syncs++; - /* Now wait for the PGRES_PIPELINE_SYNC and exit pipeline mode there */ - /* collect pending results before getting out of pipeline mode */ - return CSTATE_WAIT_RESULT; - } - - /* - * executing the expression or shell command might have taken a - * non-negligible amount of time, so reset 'now' - */ - *now = 0; - - return CSTATE_END_COMMAND; -} - - -/* - * Print log entry after completing one transaction. - * - * We print Unix-epoch timestamps in the log, so that entries can be - * correlated against other logs. - * - * XXX We could obtain the time from the caller and just shift it here, to - * avoid the cost of an extra call to pg_time_now(). - */ -static void -doLog(TState *thread, CState *st, - StatsData *agg, bool tx_skipped, double latency, double lag) -{ - FILE *logfile = thread->logfile; - pg_time_usec_t now = pg_time_now() + epoch_shift; - - Assert(use_log); - - /* - * Skip the log entry if sampling is enabled and this row doesn't belong - * to the random sample. - */ - if (sample_rate != 0.0 && - pg_prng_double(&thread->ts_sample_rs) > sample_rate) - return; - - /* should we aggregate the results or not? */ - if (agg_interval > 0) - { - pg_time_usec_t next; - - /* - * Loop until we reach the interval of the current moment, and print - * any empty intervals in between (this may happen with very low tps, - * e.g. --rate=0.1). - */ - - while ((next = agg->start_time + agg_interval * INT64CONST(1000000)) <= now) - { - double lag_sum = 0.0; - double lag_sum2 = 0.0; - double lag_min = 0.0; - double lag_max = 0.0; - int64 cnt_skipped = 0; - int64 serialization_failures = 0; - int64 deadlock_failures = 0; - int64 other_sql_failures = 0; - int64 retried = 0; - int64 retries = 0; - - /* print aggregated report to logfile */ - fprintf(logfile, INT64_FORMAT " " INT64_FORMAT " %.0f %.0f %.0f %.0f", - agg->start_time / 1000000, /* seconds since Unix epoch */ - agg->cnt, - agg->latency.sum, - agg->latency.sum2, - agg->latency.min, - agg->latency.max); + case IFSTATE_IGNORED: + case IFSTATE_ELSE_FALSE: + if (command->meta == META_IF) + conditional_stack_push(st->cstack, + IFSTATE_IGNORED); + else if (command->meta == META_ENDIF) + { + Assert(!conditional_stack_empty(st->cstack)); + conditional_stack_pop(st->cstack); + if (conditional_active(st->cstack)) + st->state = CSTATE_START_COMMAND; + } + /* could detect "else" & "elif" after "else" */ + st->command++; + break; - if (throttle_delay) - { - lag_sum = agg->lag.sum; - lag_sum2 = agg->lag.sum2; - lag_min = agg->lag.min; - lag_max = agg->lag.max; - } - fprintf(logfile, " %.0f %.0f %.0f %.0f", - lag_sum, - lag_sum2, - lag_min, - lag_max); + case IFSTATE_NONE: + case IFSTATE_TRUE: + case IFSTATE_ELSE_TRUE: + default: - if (latency_limit) - cnt_skipped = agg->cnt_skipped; - fprintf(logfile, " " INT64_FORMAT, cnt_skipped); + /* + * inconsistent if inactive, unreachable dead + * code + */ + Assert(false); + } + } + else + { + /* skip and consider next */ + st->command++; + } - if (max_tries != 1) - { - retried = agg->retried; - retries = agg->retries; - } - fprintf(logfile, " " INT64_FORMAT " " INT64_FORMAT, retried, retries); + if (st->state != CSTATE_SKIP_COMMAND) + /* out of quick skip command loop */ + break; + } + break; - if (failures_detailed) - { - serialization_failures = agg->serialization_failures; - deadlock_failures = agg->deadlock_failures; - other_sql_failures = agg->other_sql_failures; - } - fprintf(logfile, " " INT64_FORMAT " " INT64_FORMAT " " INT64_FORMAT, - serialization_failures, - deadlock_failures, - other_sql_failures); + /* + * Wait for the current SQL command to complete + */ + case CSTATE_WAIT_RESULT: + pg_log_debug("client %d receiving", st->id); - fputc('\n', logfile); + /* + * Only check for new network data if we processed all data + * fetched prior. Otherwise we end up doing a syscall for each + * individual pipelined query, which has a measurable + * performance impact. + */ + if (PQisBusy(st->con) && !PQconsumeInput(st->con)) + { + /* there's something wrong */ + commandFailed(st, "SQL", "perhaps the backend died while processing"); + st->state = CSTATE_ABORTED; + break; + } + if (PQisBusy(st->con)) + return; /* don't have the whole result yet */ - /* reset data and move to next interval */ - initStats(agg, next); - } + /* store or discard the query results */ + if (readCommandResponse(st, + sql_script[st->use_file].commands[st->command]->meta, + sql_script[st->use_file].commands[st->command]->varprefix)) + { + /* + * outside of pipeline mode: stop reading results. + * pipeline mode: continue reading results until an + * end-of-pipeline response. + */ + if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) + st->state = CSTATE_END_COMMAND; + } + else if (canRetryError(st->estatus) || canContinueOnError(st->estatus)) + st->state = CSTATE_ERROR; + else + st->state = CSTATE_ABORTED; + break; - /* accumulate the current transaction */ - accumStats(agg, tx_skipped, latency, lag, st->estatus, st->tries); - } - else - { - /* no, print raw transactions */ - if (!tx_skipped && st->estatus == ESTATUS_NO_ERROR) - fprintf(logfile, "%d " INT64_FORMAT " %.0f %d " INT64_FORMAT " " - INT64_FORMAT, - st->id, st->cnt, latency, st->use_file, - now / 1000000, now % 1000000); - else - fprintf(logfile, "%d " INT64_FORMAT " %s %d " INT64_FORMAT " " - INT64_FORMAT, - st->id, st->cnt, getResultString(tx_skipped, st->estatus), - st->use_file, now / 1000000, now % 1000000); + /* + * Wait until sleep is done. This state is entered after a + * \sleep metacommand. The behavior is similar to + * CSTATE_THROTTLE, but proceeds to CSTATE_START_COMMAND + * instead of CSTATE_START_TX. + */ + case CSTATE_SLEEP: + pg_time_now_lazy(&now); + if (now < st->sleep_until) + return; /* still sleeping, nothing to do here */ + /* Else done sleeping. */ + st->state = CSTATE_END_COMMAND; + break; - if (throttle_delay) - fprintf(logfile, " %.0f", lag); - if (max_tries != 1) - fprintf(logfile, " %u", st->tries - 1); - fputc('\n', logfile); - } -} + /* + * End of command: record stats and proceed to next command. + */ + case CSTATE_END_COMMAND: -/* - * Accumulate and report statistics at end of a transaction. - * - * (This is also called when a transaction is late and thus skipped. - * Note that even skipped and failed transactions are counted in the CState - * "cnt" field.) - */ -static void -processXactStats(TState *thread, CState *st, pg_time_usec_t *now, - bool tx_skipped, StatsData *agg) -{ - double latency = 0.0, - lag = 0.0; - bool detailed = progress || throttle_delay || latency_limit || - use_log || per_script_stats; + /* + * command completed: accumulate per-command execution times + * in thread-local data structure, if per-command latencies + * are requested. + */ + if (report_per_command) + { + pg_time_now_lazy(&now); - if (detailed && !tx_skipped && st->estatus == ESTATUS_NO_ERROR) - { - pg_time_now_lazy(now); + command = sql_script[st->use_file].commands[st->command]; + /* XXX could use a mutex here, but we choose not to */ + addToSimpleStats(&command->stats, + PG_TIME_GET_DOUBLE(now - st->stmt_begin)); + } - /* compute latency & lag */ - latency = (*now) - st->txn_scheduled; - lag = st->txn_begin - st->txn_scheduled; - } + /* Go ahead with next command, to be executed or skipped */ + st->command++; + st->state = conditional_active(st->cstack) ? + CSTATE_START_COMMAND : CSTATE_SKIP_COMMAND; + break; - /* keep detailed thread stats */ - accumStats(&thread->stats, tx_skipped, latency, lag, st->estatus, st->tries); + /* + * Clean up after an error. + */ + case CSTATE_ERROR: + { + TStatus tstatus; - /* count transactions over the latency limit, if needed */ - if (latency_limit && latency > latency_limit) - thread->latency_late++; + Assert(st->estatus != ESTATUS_NO_ERROR); - /* client stat is just counting */ - st->cnt++; + /* Clear the conditional stack */ + conditional_stack_reset(st->cstack); - if (use_log) - doLog(thread, st, agg, tx_skipped, latency, lag); + /* Read and discard until a sync point in pipeline mode */ + if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) + { + if (!discardUntilSync(st)) + { + st->state = CSTATE_ABORTED; + break; + } + } - /* XXX could use a mutex here, but we choose not to */ - if (per_script_stats) - accumStats(&sql_script[st->use_file].stats, tx_skipped, latency, lag, - st->estatus, st->tries); -} + /* + * Check if we have a (failed) transaction block or not, + * and roll it back if any. + */ + tstatus = getTransactionStatus(st->con); + if (tstatus == TSTATUS_IN_BLOCK) + { + /* Try to rollback a (failed) transaction block. */ + if (!PQsendQuery(st->con, "ROLLBACK")) + { + pg_log_error("client %d aborted: failed to send sql command for rolling back the failed transaction", + st->id); + st->state = CSTATE_ABORTED; + } + else + st->state = CSTATE_WAIT_ROLLBACK_RESULT; + } + else if (tstatus == TSTATUS_IDLE) + { + /* + * If time is over, we're done; otherwise, check if we + * can retry the error. + */ + st->state = timer_exceeded ? CSTATE_FINISHED : + doRetry(st, &now) ? CSTATE_RETRY : CSTATE_FAILURE; + } + else + { + if (tstatus == TSTATUS_CONN_ERROR) + pg_log_error("perhaps the backend died while processing"); + pg_log_error("client %d aborted while receiving the transaction status", st->id); + st->state = CSTATE_ABORTED; + } + break; + } -/* discard connections */ -static void -disconnect_all(CState *state, int length) -{ - int i; + /* + * Wait for the rollback command to complete + */ + case CSTATE_WAIT_ROLLBACK_RESULT: + { + PGresult *res; - for (i = 0; i < length; i++) - finishCon(&state[i]); -} + pg_log_debug("client %d receiving", st->id); + if (!PQconsumeInput(st->con)) + { + pg_log_error("client %d aborted while rolling back the transaction after an error; perhaps the backend died while processing", + st->id); + st->state = CSTATE_ABORTED; + break; + } + if (PQisBusy(st->con)) + return; /* don't have the whole result yet */ -/* - * Remove old pgbench tables, if any exist - */ -static void -initDropTables(PGconn *con) -{ - fprintf(stderr, "dropping old tables...\n"); + /* + * Read and discard the query result; + */ + res = PQgetResult(st->con); + switch (PQresultStatus(res)) + { + case PGRES_COMMAND_OK: + /* OK */ + PQclear(res); + /* null must be returned */ + res = PQgetResult(st->con); + Assert(res == NULL); - /* - * We drop all the tables in one command, so that whether there are - * foreign key dependencies or not doesn't matter. - */ - executeStatement(con, "drop table if exists " - "pgbench_accounts, " - "pgbench_branches, " - "pgbench_history, " - "pgbench_tellers"); -} + /* + * If time is over, we're done; otherwise, check + * if we can retry the error. + */ + st->state = timer_exceeded ? CSTATE_FINISHED : + doRetry(st, &now) ? CSTATE_RETRY : CSTATE_FAILURE; + break; + default: + pg_log_error("client %d aborted while rolling back the transaction after an error; %s", + st->id, PQerrorMessage(st->con)); + PQclear(res); + st->state = CSTATE_ABORTED; + break; + } + break; + } -/* - * Create "pgbench_accounts" partitions if needed. - * - * This is the larger table of pgbench default tpc-b like schema - * with a known size, so we choose to partition it. - */ -static void -createPartitions(PGconn *con) -{ - PQExpBufferData query; + /* + * Retry the transaction after an error. + */ + case CSTATE_RETRY: + command = sql_script[st->use_file].commands[st->command]; - /* we must have to create some partitions */ - Assert(partitions > 0); + /* + * Inform that the transaction will be retried after the + * error. + */ + if (verbose_errors) + printVerboseErrorMessages(st, &now, true); - fprintf(stderr, "creating %d partitions...\n", partitions); + /* Count tries and retries */ + st->tries++; + command->retries++; - initPQExpBuffer(&query); + /* + * Reset the random state as they were at the beginning of the + * transaction. + */ + st->cs_func_rs = st->random_state; - for (int p = 1; p <= partitions; p++) - { - if (partition_method == PART_RANGE) - { - int64 part_size = (naccounts * (int64) scale + partitions - 1) / partitions; + /* Process the first transaction command. */ + st->command = 0; + st->estatus = ESTATUS_NO_ERROR; + st->state = CSTATE_START_COMMAND; + break; - printfPQExpBuffer(&query, - "create%s table pgbench_accounts_%d\n" - " partition of pgbench_accounts\n" - " for values from (", - unlogged_tables ? " unlogged" : "", p); + /* + * Record a failed transaction. + */ + case CSTATE_FAILURE: + command = sql_script[st->use_file].commands[st->command]; - /* - * For RANGE, we use open-ended partitions at the beginning and - * end to allow any valid value for the primary key. Although the - * actual minimum and maximum values can be derived from the - * scale, it is more generic and the performance is better. - */ - if (p == 1) - appendPQExpBufferStr(&query, "minvalue"); - else - appendPQExpBuffer(&query, INT64_FORMAT, (p - 1) * part_size + 1); + /* Accumulate the failure. */ + command->failures++; - appendPQExpBufferStr(&query, ") to ("); + /* + * Inform that the failed transaction will not be retried. + */ + if (verbose_errors) + printVerboseErrorMessages(st, &now, false); - if (p < partitions) - appendPQExpBuffer(&query, INT64_FORMAT, p * part_size + 1); - else - appendPQExpBufferStr(&query, "maxvalue"); + /* End the failed transaction. */ + st->state = CSTATE_END_TX; + break; - appendPQExpBufferChar(&query, ')'); - } - else if (partition_method == PART_HASH) - printfPQExpBuffer(&query, - "create%s table pgbench_accounts_%d\n" - " partition of pgbench_accounts\n" - " for values with (modulus %d, remainder %d)", - unlogged_tables ? " unlogged" : "", p, - partitions, p - 1); - else /* cannot get there */ - Assert(0); + /* + * End of transaction (end of script, really). + */ + case CSTATE_END_TX: + { + TStatus tstatus; - /* - * Per ddlinfo in initCreateTables, fillfactor is needed on table - * pgbench_accounts. - */ - appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor); + /* transaction finished: calculate latency and do log */ + processXactStats(thread, st, &now, false, agg); - executeStatement(con, query.data); - } + /* + * missing \endif... cannot happen if CheckConditional was + * okay + */ + Assert(conditional_stack_empty(st->cstack)); - termPQExpBuffer(&query); -} + /* + * We must complete all the transaction blocks that were + * started in this script. + */ + tstatus = getTransactionStatus(st->con); + if (tstatus == TSTATUS_IN_BLOCK) + { + pg_log_error("client %d aborted: end of script reached without completing the last transaction", + st->id); + st->state = CSTATE_ABORTED; + break; + } + else if (tstatus != TSTATUS_IDLE) + { + if (tstatus == TSTATUS_CONN_ERROR) + pg_log_error("perhaps the backend died while processing"); -/* - * Create pgbench's standard tables - */ -static void -initCreateTables(PGconn *con) -{ - /* - * Note: TPC-B requires at least 100 bytes per row, and the "filler" - * fields in these table declarations were intended to comply with that. - * The pgbench_accounts table complies with that because the "filler" - * column is set to blank-padded empty string. But for all other tables - * the columns default to NULL and so don't actually take any space. We - * could fix that by giving them non-null default values. However, that - * would completely break comparability of pgbench results with prior - * versions. Since pgbench has never pretended to be fully TPC-B compliant - * anyway, we stick with the historical behavior. - */ - struct ddlinfo - { - const char *table; /* table name */ - const char *smcols; /* column decls if accountIDs are 32 bits */ - const char *bigcols; /* column decls if accountIDs are 64 bits */ - int declare_fillfactor; - }; - static const struct ddlinfo DDLs[] = { - { - "pgbench_history", - "tid int,bid int,aid int,delta int,mtime timestamp,filler char(22)", - "tid int,bid int,aid bigint,delta int,mtime timestamp,filler char(22)", - 0 - }, - { - "pgbench_tellers", - "tid int not null,bid int,tbalance int,filler char(84)", - "tid int not null,bid int,tbalance int,filler char(84)", - 1 - }, - { - "pgbench_accounts", - "aid int not null,bid int,abalance int,filler char(84)", - "aid bigint not null,bid int,abalance int,filler char(84)", - 1 - }, - { - "pgbench_branches", - "bid int not null,bbalance int,filler char(88)", - "bid int not null,bbalance int,filler char(88)", - 1 - } - }; - PQExpBufferData query; + pg_log_error("client %d aborted while receiving the transaction status", st->id); + st->state = CSTATE_ABORTED; + break; + } - fprintf(stderr, "creating tables...\n"); + if (is_connect) + { + pg_time_usec_t start = now; - initPQExpBuffer(&query); + pg_time_now_lazy(&start); + finishCon(st); + now = pg_time_now(); + thread->conn_duration += now - start; + } - for (size_t i = 0; i < lengthof(DDLs); i++) - { - const struct ddlinfo *ddl = &DDLs[i]; + if ((st->cnt >= nxacts && duration <= 0) || timer_exceeded) + { + /* script completed */ + st->state = CSTATE_FINISHED; + break; + } - /* Construct new create table statement. */ - printfPQExpBuffer(&query, "create%s table %s(%s)", - (unlogged_tables && partition_method == PART_NONE) ? " unlogged" : "", - ddl->table, - (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols); + /* next transaction (script) */ + st->state = CSTATE_CHOOSE_SCRIPT; - /* Partition pgbench_accounts table */ - if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0) - appendPQExpBuffer(&query, - " partition by %s (aid)", PARTITION_METHOD[partition_method]); - else if (ddl->declare_fillfactor) - { - /* fillfactor is only expected on actual tables */ - appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor); - } + /* + * Ensure that we always return on this point, so as to + * avoid an infinite loop if the script only contains meta + * commands. + */ + return; + } - if (tablespace != NULL) - { - char *escape_tablespace; + /* + * Final states. Close the connection if it's still open. + */ + case CSTATE_ABORTED: + case CSTATE_FINISHED: - escape_tablespace = PQescapeIdentifier(con, tablespace, strlen(tablespace)); - appendPQExpBuffer(&query, " tablespace %s", escape_tablespace); - PQfreemem(escape_tablespace); + /* + * Don't measure the disconnection delays here even if in + * CSTATE_FINISHED and -C/--connect option is specified. + * Because in this case all the connections that this thread + * established are closed at the end of transactions and the + * disconnection delays should have already been measured at + * that moment. + * + * In CSTATE_ABORTED state, the measurement is no longer + * necessary because we cannot report complete results anyways + * in this case. + */ + finishCon(st); + return; } - - executeStatement(con, query.data); } - - termPQExpBuffer(&query); - - if (partition_method != PART_NONE) - createPartitions(con); } /* - * Truncate away any old data, in one command in case there are foreign keys + * Subroutine for advanceConnectionState -- initiate or execute the current + * meta command, and return the next state to set. + * + * *now is updated to the current time, unless the command is expected to + * take no time to execute. */ -static void -initTruncateTables(PGconn *con) +static ConnectionStateEnum +executeMetaCommand(CState *st, pg_time_usec_t *now) { - executeStatement(con, "truncate table " - "pgbench_accounts, " - "pgbench_branches, " - "pgbench_history, " - "pgbench_tellers"); -} + Command *command = sql_script[st->use_file].commands[st->command]; + int argc; + char **argv; -static void -initBranch(PQExpBufferData *sql, int64 curr) -{ - /* "filler" column uses NULL */ - printfPQExpBuffer(sql, - INT64_FORMAT "\t0\t\\N\n", - curr + 1); -} + Assert(command != NULL && command->type == META_COMMAND); -static void -initTeller(PQExpBufferData *sql, int64 curr) -{ - /* "filler" column uses NULL */ - printfPQExpBuffer(sql, - INT64_FORMAT "\t" INT64_FORMAT "\t0\t\\N\n", - curr + 1, curr / ntellers + 1); -} + argc = command->argc; + argv = command->argv; -static void -initAccount(PQExpBufferData *sql, int64 curr) -{ - /* "filler" column defaults to blank padded empty string */ - printfPQExpBuffer(sql, - INT64_FORMAT "\t" INT64_FORMAT "\t0\t\n", - curr + 1, curr / naccounts + 1); -} + if (unlikely(__pg_log_level <= PG_LOG_DEBUG)) + { + PQExpBufferData buf; -static void -initPopulateTable(PGconn *con, const char *table, int64 base, - initRowMethod init_row) -{ - int n; - int64 k; - int chars = 0; - int prev_chars = 0; - PGresult *res; - PQExpBufferData sql; - char copy_statement[256]; - const char *copy_statement_fmt = "copy %s from stdin"; - int64 total = base * scale; + initPQExpBuffer(&buf); - /* used to track elapsed time and estimate of the remaining time */ - pg_time_usec_t start; - int log_interval = 1; + printfPQExpBuffer(&buf, "client %d executing \\%s", st->id, argv[0]); + for (int i = 1; i < argc; i++) + appendPQExpBuffer(&buf, " %s", argv[i]); - /* Stay on the same line if reporting to a terminal */ - char eol = isatty(fileno(stderr)) ? '\r' : '\n'; + pg_log_debug("%s", buf.data); - initPQExpBuffer(&sql); + termPQExpBuffer(&buf); + } - /* Use COPY with FREEZE on v14 and later for all ordinary tables */ - if ((PQserverVersion(con) >= 140000) && - get_table_relkind(con, table) == RELKIND_RELATION) - copy_statement_fmt = "copy %s from stdin with (freeze on)"; + if (command->meta == META_SLEEP) + { + int usec; + /* + * A \sleep doesn't execute anything, we just get the delay from the + * argument, and enter the CSTATE_SLEEP state. (The per-command + * latency will be recorded in CSTATE_SLEEP state, not here, after the + * delay has elapsed.) + */ + if (!evaluateSleep(&st->variables, argc, argv, &usec)) + { + commandFailed(st, "sleep", "execution of meta-command failed"); + return CSTATE_ABORTED; + } - n = pg_snprintf(copy_statement, sizeof(copy_statement), copy_statement_fmt, table); - if (n >= sizeof(copy_statement)) - pg_fatal("invalid buffer size: must be at least %d characters long", n); - else if (n == -1) - pg_fatal("invalid format string"); + pg_time_now_lazy(now); + st->sleep_until = (*now) + usec; + return CSTATE_SLEEP; + } + else if (command->meta == META_SET) + { + PgBenchExpr *expr = command->expr; + PgBenchValue result; - res = PQexec(con, copy_statement); + if (!evaluateExpr(st, expr, &result)) + { + commandFailed(st, argv[0], "evaluation of meta-command failed"); + return CSTATE_ABORTED; + } - if (PQresultStatus(res) != PGRES_COPY_IN) - pg_fatal("unexpected copy in result: %s", PQerrorMessage(con)); - PQclear(res); + if (!putVariableValue(&st->variables, argv[0], argv[1], &result)) + { + commandFailed(st, "set", "assignment of meta-command failed"); + return CSTATE_ABORTED; + } + } + else if (command->meta == META_IF) + { + /* backslash commands with an expression to evaluate */ + PgBenchExpr *expr = command->expr; + PgBenchValue result; + bool cond; - start = pg_time_now(); + if (!evaluateExpr(st, expr, &result)) + { + commandFailed(st, argv[0], "evaluation of meta-command failed"); + return CSTATE_ABORTED; + } - for (k = 0; k < total; k++) + cond = valueTruth(&result); + conditional_stack_push(st->cstack, cond ? IFSTATE_TRUE : IFSTATE_FALSE); + } + else if (command->meta == META_ELIF) + { + /* backslash commands with an expression to evaluate */ + PgBenchExpr *expr = command->expr; + PgBenchValue result; + bool cond; + + if (conditional_stack_peek(st->cstack) == IFSTATE_TRUE) + { + /* elif after executed block, skip eval and wait for endif. */ + conditional_stack_poke(st->cstack, IFSTATE_IGNORED); + return CSTATE_END_COMMAND; + } + + if (!evaluateExpr(st, expr, &result)) + { + commandFailed(st, argv[0], "evaluation of meta-command failed"); + return CSTATE_ABORTED; + } + + cond = valueTruth(&result); + Assert(conditional_stack_peek(st->cstack) == IFSTATE_FALSE); + conditional_stack_poke(st->cstack, cond ? IFSTATE_TRUE : IFSTATE_FALSE); + } + else if (command->meta == META_ELSE) + { + switch (conditional_stack_peek(st->cstack)) + { + case IFSTATE_TRUE: + conditional_stack_poke(st->cstack, IFSTATE_ELSE_FALSE); + break; + case IFSTATE_FALSE: /* inconsistent if active */ + case IFSTATE_IGNORED: /* inconsistent if active */ + case IFSTATE_NONE: /* else without if */ + case IFSTATE_ELSE_TRUE: /* else after else */ + case IFSTATE_ELSE_FALSE: /* else after else */ + default: + /* dead code if conditional check is ok */ + Assert(false); + } + } + else if (command->meta == META_ENDIF) + { + Assert(!conditional_stack_empty(st->cstack)); + conditional_stack_pop(st->cstack); + } + else if (command->meta == META_SETSHELL) + { + if (!runShellCommand(&st->variables, argv[1], argv + 2, argc - 2)) + { + commandFailed(st, "setshell", "execution of meta-command failed"); + return CSTATE_ABORTED; + } + } + else if (command->meta == META_SHELL) + { + if (!runShellCommand(&st->variables, NULL, argv + 1, argc - 1)) + { + commandFailed(st, "shell", "execution of meta-command failed"); + return CSTATE_ABORTED; + } + } + else if (command->meta == META_STARTPIPELINE) { - int64 j = k + 1; - - init_row(&sql, k); - if (PQputline(con, sql.data)) - pg_fatal("PQputline failed"); - - if (CancelRequested) - break; - /* - * If we want to stick with the original logging, print a message each - * 100k inserted rows. + * In pipeline mode, we use a workflow based on libpq pipeline + * functions. */ - if ((!use_quiet) && (j % 100000 == 0)) + if (querymode == QUERY_SIMPLE) { - double elapsed_sec = PG_TIME_GET_DOUBLE(pg_time_now() - start); - double remaining_sec = ((double) total - j) * elapsed_sec / j; + commandFailed(st, "startpipeline", "cannot use pipeline mode with the simple query protocol"); + return CSTATE_ABORTED; + } - chars = fprintf(stderr, INT64_FORMAT " of " INT64_FORMAT " tuples (%d%%) of %s done (elapsed %.2f s, remaining %.2f s)", - j, total, - (int) ((j * 100) / total), - table, elapsed_sec, remaining_sec); + /* + * If we're in prepared-query mode, we need to prepare all the + * commands that are inside the pipeline before we actually start the + * pipeline itself. This solves the problem that running BEGIN + * ISOLATION LEVEL SERIALIZABLE in a pipeline would fail due to a + * snapshot having been acquired by the prepare within the pipeline. + */ + if (querymode == QUERY_PREPARED) + prepareCommandsInPipeline(st); - /* - * If the previous progress message is longer than the current - * one, add spaces to the current line to fully overwrite any - * remaining characters from the previous message. - */ - if (prev_chars > chars) - fprintf(stderr, "%*c", prev_chars - chars, ' '); - fputc(eol, stderr); - prev_chars = chars; + if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) + { + commandFailed(st, "startpipeline", "already in pipeline mode"); + return CSTATE_ABORTED; } - /* let's not call the timing for each row, but only each 100 rows */ - else if (use_quiet && (j % 100 == 0)) + if (PQenterPipelineMode(st->con) == 0) { - double elapsed_sec = PG_TIME_GET_DOUBLE(pg_time_now() - start); - double remaining_sec = ((double) total - j) * elapsed_sec / j; - - /* have we reached the next interval (or end)? */ - if ((j == total) || (elapsed_sec >= log_interval * LOG_STEP_SECONDS)) - { - chars = fprintf(stderr, INT64_FORMAT " of " INT64_FORMAT " tuples (%d%%) of %s done (elapsed %.2f s, remaining %.2f s)", - j, total, - (int) ((j * 100) / total), - table, elapsed_sec, remaining_sec); - - /* - * If the previous progress message is longer than the current - * one, add spaces to the current line to fully overwrite any - * remaining characters from the previous message. - */ - if (prev_chars > chars) - fprintf(stderr, "%*c", prev_chars - chars, ' '); - fputc(eol, stderr); - prev_chars = chars; - - /* skip to the next interval */ - log_interval = (int) ceil(elapsed_sec / LOG_STEP_SECONDS); - } + commandFailed(st, "startpipeline", "failed to enter pipeline mode"); + return CSTATE_ABORTED; } } + else if (command->meta == META_SYNCPIPELINE) + { + if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) + { + commandFailed(st, "syncpipeline", "not in pipeline mode"); + return CSTATE_ABORTED; + } + if (PQsendPipelineSync(st->con) == 0) + { + commandFailed(st, "syncpipeline", "failed to send a pipeline sync"); + return CSTATE_ABORTED; + } + st->num_syncs++; + } + else if (command->meta == META_ENDPIPELINE) + { + if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) + { + commandFailed(st, "endpipeline", "not in pipeline mode"); + return CSTATE_ABORTED; + } + if (!PQpipelineSync(st->con)) + { + commandFailed(st, "endpipeline", "failed to send a pipeline sync"); + return CSTATE_ABORTED; + } + st->num_syncs++; + /* Now wait for the PGRES_PIPELINE_SYNC and exit pipeline mode there */ + /* collect pending results before getting out of pipeline mode */ + return CSTATE_WAIT_RESULT; + } - if (chars != 0 && eol != '\n') - fprintf(stderr, "%*c\r", chars, ' '); /* Clear the current line */ - - if (PQputline(con, "\\.\n")) - pg_fatal("very last PQputline failed"); - if (PQendcopy(con)) - pg_fatal("PQendcopy failed"); + /* + * executing the expression or shell command might have taken a + * non-negligible amount of time, so reset 'now' + */ + *now = 0; - termPQExpBuffer(&sql); + return CSTATE_END_COMMAND; } + /* - * Fill the standard tables with some data generated and sent from the client. + * Print log entry after completing one transaction. * - * The filler column is NULL in pgbench_branches and pgbench_tellers, and is - * a blank-padded string in pgbench_accounts. + * We print Unix-epoch timestamps in the log, so that entries can be + * correlated against other logs. + * + * XXX We could obtain the time from the caller and just shift it here, to + * avoid the cost of an extra call to pg_time_now(). */ static void -initGenerateDataClientSide(PGconn *con) +doLog(TState *thread, CState *st, + StatsData *agg, bool tx_skipped, double latency, double lag) { - fprintf(stderr, "generating data (client-side)...\n"); - - /* - * we do all of this in one transaction to enable the backend's - * data-loading optimizations - */ - executeStatement(con, "begin"); + FILE *logfile = thread->logfile; + pg_time_usec_t now = pg_time_now() + epoch_shift; - /* truncate away any old data */ - initTruncateTables(con); + Assert(use_log); /* - * fill branches, tellers, accounts in that order in case foreign keys - * already exist + * Skip the log entry if sampling is enabled and this row doesn't belong + * to the random sample. */ - initPopulateTable(con, "pgbench_branches", nbranches, initBranch); - initPopulateTable(con, "pgbench_tellers", ntellers, initTeller); - initPopulateTable(con, "pgbench_accounts", naccounts, initAccount); - - executeStatement(con, "commit"); -} + if (sample_rate != 0.0 && + pg_prng_double(&thread->ts_sample_rs) > sample_rate) + return; -/* - * Fill the standard tables with some data generated on the server - * - * As already the case with the client-side data generation, the filler - * column defaults to NULL in pgbench_branches and pgbench_tellers, - * and is a blank-padded string in pgbench_accounts. - */ -static void -initGenerateDataServerSide(PGconn *con) -{ - PQExpBufferData sql; + /* should we aggregate the results or not? */ + if (agg_interval > 0) + { + pg_time_usec_t next; - fprintf(stderr, "generating data (server-side)...\n"); + /* + * Loop until we reach the interval of the current moment, and print + * any empty intervals in between (this may happen with very low tps, + * e.g. --rate=0.1). + */ - /* - * we do all of this in one transaction to enable the backend's - * data-loading optimizations - */ - executeStatement(con, "begin"); + while ((next = agg->start_time + agg_interval * INT64CONST(1000000)) <= now) + { + double lag_sum = 0.0; + double lag_sum2 = 0.0; + double lag_min = 0.0; + double lag_max = 0.0; + int64 cnt_skipped = 0; + int64 serialization_failures = 0; + int64 deadlock_failures = 0; + int64 other_sql_failures = 0; + int64 retried = 0; + int64 retries = 0; - /* truncate away any old data */ - initTruncateTables(con); + /* print aggregated report to logfile */ + fprintf(logfile, INT64_FORMAT " " INT64_FORMAT " %.0f %.0f %.0f %.0f", + agg->start_time / 1000000, /* seconds since Unix epoch */ + agg->cnt, + agg->latency.sum, + agg->latency.sum2, + agg->latency.min, + agg->latency.max); - initPQExpBuffer(&sql); + if (throttle_delay) + { + lag_sum = agg->lag.sum; + lag_sum2 = agg->lag.sum2; + lag_min = agg->lag.min; + lag_max = agg->lag.max; + } + fprintf(logfile, " %.0f %.0f %.0f %.0f", + lag_sum, + lag_sum2, + lag_min, + lag_max); - printfPQExpBuffer(&sql, - "insert into pgbench_branches(bid,bbalance) " - "select bid, 0 " - "from generate_series(1, %d) as bid", nbranches * scale); - executeStatement(con, sql.data); + if (latency_limit) + cnt_skipped = agg->cnt_skipped; + fprintf(logfile, " " INT64_FORMAT, cnt_skipped); - printfPQExpBuffer(&sql, - "insert into pgbench_tellers(tid,bid,tbalance) " - "select tid, (tid - 1) / %d + 1, 0 " - "from generate_series(1, %d) as tid", ntellers, ntellers * scale); - executeStatement(con, sql.data); + if (max_tries != 1) + { + retried = agg->retried; + retries = agg->retries; + } + fprintf(logfile, " " INT64_FORMAT " " INT64_FORMAT, retried, retries); - printfPQExpBuffer(&sql, - "insert into pgbench_accounts(aid,bid,abalance,filler) " - "select aid, (aid - 1) / %d + 1, 0, '' " - "from generate_series(1, " INT64_FORMAT ") as aid", - naccounts, (int64) naccounts * scale); - executeStatement(con, sql.data); + if (failures_detailed) + { + serialization_failures = agg->serialization_failures; + deadlock_failures = agg->deadlock_failures; + other_sql_failures = agg->other_sql_failures; + } + fprintf(logfile, " " INT64_FORMAT " " INT64_FORMAT " " INT64_FORMAT, + serialization_failures, + deadlock_failures, + other_sql_failures); - termPQExpBuffer(&sql); + fputc('\n', logfile); - executeStatement(con, "commit"); -} + /* reset data and move to next interval */ + initStats(agg, next); + } -/* - * Invoke vacuum on the standard tables - */ -static void -initVacuum(PGconn *con) -{ - fprintf(stderr, "vacuuming...\n"); - executeStatement(con, "vacuum analyze pgbench_branches"); - executeStatement(con, "vacuum analyze pgbench_tellers"); - executeStatement(con, "vacuum analyze pgbench_accounts"); - executeStatement(con, "vacuum analyze pgbench_history"); + /* accumulate the current transaction */ + accumStats(agg, tx_skipped, latency, lag, st->estatus, st->tries); + } + else + { + /* no, print raw transactions */ + if (!tx_skipped && st->estatus == ESTATUS_NO_ERROR) + fprintf(logfile, "%d " INT64_FORMAT " %.0f %d " INT64_FORMAT " " + INT64_FORMAT, + st->id, st->cnt, latency, st->use_file, + now / 1000000, now % 1000000); + else + fprintf(logfile, "%d " INT64_FORMAT " %s %d " INT64_FORMAT " " + INT64_FORMAT, + st->id, st->cnt, getResultString(tx_skipped, st->estatus), + st->use_file, now / 1000000, now % 1000000); + + if (throttle_delay) + fprintf(logfile, " %.0f", lag); + if (max_tries != 1) + fprintf(logfile, " %u", st->tries - 1); + fputc('\n', logfile); + } } /* - * Create primary keys on the standard tables + * Accumulate and report statistics at end of a transaction. + * + * (This is also called when a transaction is late and thus skipped. + * Note that even skipped and failed transactions are counted in the CState + * "cnt" field.) */ static void -initCreatePKeys(PGconn *con) +processXactStats(TState *thread, CState *st, pg_time_usec_t *now, + bool tx_skipped, StatsData *agg) { - static const char *const DDLINDEXes[] = { - "alter table pgbench_branches add primary key (bid)", - "alter table pgbench_tellers add primary key (tid)", - "alter table pgbench_accounts add primary key (aid)" - }; - PQExpBufferData query; - - fprintf(stderr, "creating primary keys...\n"); - initPQExpBuffer(&query); + double latency = 0.0, + lag = 0.0; + bool detailed = progress || throttle_delay || latency_limit || + use_log || per_script_stats; - for (size_t i = 0; i < lengthof(DDLINDEXes); i++) + if (detailed && !tx_skipped && st->estatus == ESTATUS_NO_ERROR) { - resetPQExpBuffer(&query); - appendPQExpBufferStr(&query, DDLINDEXes[i]); + pg_time_now_lazy(now); - if (index_tablespace != NULL) - { - char *escape_tablespace; + /* compute latency & lag */ + latency = (*now) - st->txn_scheduled; + lag = st->txn_begin - st->txn_scheduled; + } - escape_tablespace = PQescapeIdentifier(con, index_tablespace, - strlen(index_tablespace)); - appendPQExpBuffer(&query, " using index tablespace %s", escape_tablespace); - PQfreemem(escape_tablespace); - } + /* keep detailed thread stats */ + accumStats(&thread->stats, tx_skipped, latency, lag, st->estatus, st->tries); - executeStatement(con, query.data); - } + /* count transactions over the latency limit, if needed */ + if (latency_limit && latency > latency_limit) + thread->latency_late++; - termPQExpBuffer(&query); + /* client stat is just counting */ + st->cnt++; + + if (use_log) + doLog(thread, st, agg, tx_skipped, latency, lag); + + /* XXX could use a mutex here, but we choose not to */ + if (per_script_stats) + accumStats(&sql_script[st->use_file].stats, tx_skipped, latency, lag, + st->estatus, st->tries); } -/* - * Create foreign key constraints between the standard tables - */ + +/* discard connections */ static void -initCreateFKeys(PGconn *con) +disconnect_all(CState *state, int length) { - static const char *const DDLKEYs[] = { - "alter table pgbench_tellers add constraint pgbench_tellers_bid_fkey foreign key (bid) references pgbench_branches", - "alter table pgbench_accounts add constraint pgbench_accounts_bid_fkey foreign key (bid) references pgbench_branches", - "alter table pgbench_history add constraint pgbench_history_bid_fkey foreign key (bid) references pgbench_branches", - "alter table pgbench_history add constraint pgbench_history_tid_fkey foreign key (tid) references pgbench_tellers", - "alter table pgbench_history add constraint pgbench_history_aid_fkey foreign key (aid) references pgbench_accounts" - }; + int i; - fprintf(stderr, "creating foreign keys...\n"); - for (size_t i = 0; i < lengthof(DDLKEYs); i++) - { - executeStatement(con, DDLKEYs[i]); - } + for (i = 0; i < length; i++) + finishCon(&state[i]); } /* - * Validate an initialization-steps string - * - * (We could just leave it to runInitSteps() to fail if there are wrong - * characters, but since initialization can take awhile, it seems friendlier - * to check during option parsing.) + * Remove old pgbench tables, if any exist */ static void -checkInitSteps(const char *initialize_steps) +initDropTables(PGconn *con) { - if (initialize_steps[0] == '\0') - pg_fatal("no initialization steps specified"); + fprintf(stderr, "dropping old tables...\n"); - for (const char *step = initialize_steps; *step != '\0'; step++) - { - if (strchr(ALL_INIT_STEPS " ", *step) == NULL) - { - pg_log_error("unrecognized initialization step \"%c\"", *step); - pg_log_error_detail("Allowed step characters are: \"" ALL_INIT_STEPS "\"."); - exit(1); - } - } + /* + * We drop all the tables in one command, so that whether there are + * foreign key dependencies or not doesn't matter. + */ + executeStatement(con, "drop table if exists " + "pgbench_accounts, " + "pgbench_branches, " + "pgbench_history, " + "pgbench_tellers"); } /* - * Invoke each initialization step in the given string + * Create "pgbench_accounts" partitions if needed. + * + * This is the larger table of pgbench default tpc-b like schema + * with a known size, so we choose to partition it. */ static void -runInitSteps(const char *initialize_steps) +createPartitions(PGconn *con) { - PQExpBufferData stats; - PGconn *con; - const char *step; - double run_time = 0.0; - bool first = true; + PQExpBufferData query; - initPQExpBuffer(&stats); + /* we must have to create some partitions */ + Assert(partitions > 0); - if ((con = doConnect()) == NULL) - pg_fatal("could not create connection for initialization"); + fprintf(stderr, "creating %d partitions...\n", partitions); - setup_cancel_handler(NULL); - SetCancelConn(con); + initPQExpBuffer(&query); - for (step = initialize_steps; *step != '\0'; step++) + for (int p = 1; p <= partitions; p++) { - char *op = NULL; - pg_time_usec_t start = pg_time_now(); - - switch (*step) + if (partition_method == PART_RANGE) { - case 'd': - op = "drop tables"; - initDropTables(con); - break; - case 't': - op = "create tables"; - initCreateTables(con); - break; - case 'g': - op = "client-side generate"; - initGenerateDataClientSide(con); - break; - case 'G': - op = "server-side generate"; - initGenerateDataServerSide(con); - break; - case 'v': - op = "vacuum"; - initVacuum(con); - break; - case 'p': - op = "primary keys"; - initCreatePKeys(con); - break; - case 'f': - op = "foreign keys"; - initCreateFKeys(con); - break; - case ' ': - break; /* ignore */ - default: - pg_log_error("unrecognized initialization step \"%c\"", *step); - PQfinish(con); - exit(1); - } + int64 part_size = (naccounts * (int64) scale + partitions - 1) / partitions; - if (op != NULL) - { - double elapsed_sec = PG_TIME_GET_DOUBLE(pg_time_now() - start); + printfPQExpBuffer(&query, + "create%s table pgbench_accounts_%d\n" + " partition of pgbench_accounts\n" + " for values from (", + unlogged_tables ? " unlogged" : "", p); - if (!first) - appendPQExpBufferStr(&stats, ", "); + /* + * For RANGE, we use open-ended partitions at the beginning and + * end to allow any valid value for the primary key. Although the + * actual minimum and maximum values can be derived from the + * scale, it is more generic and the performance is better. + */ + if (p == 1) + appendPQExpBufferStr(&query, "minvalue"); else - first = false; + appendPQExpBuffer(&query, INT64_FORMAT, (p - 1) * part_size + 1); - appendPQExpBuffer(&stats, "%s %.2f s", op, elapsed_sec); + appendPQExpBufferStr(&query, ") to ("); + + if (p < partitions) + appendPQExpBuffer(&query, INT64_FORMAT, p * part_size + 1); + else + appendPQExpBufferStr(&query, "maxvalue"); + + appendPQExpBufferChar(&query, ')'); + } + else if (partition_method == PART_HASH) + printfPQExpBuffer(&query, + "create%s table pgbench_accounts_%d\n" + " partition of pgbench_accounts\n" + " for values with (modulus %d, remainder %d)", + unlogged_tables ? " unlogged" : "", p, + partitions, p - 1); + else /* cannot get there */ + Assert(0); + + /* + * Per ddlinfo in initCreateTables, fillfactor is needed on table + * pgbench_accounts. + */ + appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor); - run_time += elapsed_sec; - } + executeStatement(con, query.data); } - fprintf(stderr, "done in %.2f s (%s).\n", run_time, stats.data); - ResetCancelConn(); - PQfinish(con); - termPQExpBuffer(&stats); + termPQExpBuffer(&query); } /* - * Extract pgbench table information into global variables scale, - * partition_method and partitions. + * Create pgbench's standard tables */ static void -GetTableInfo(PGconn *con, bool scale_given) +initCreateTables(PGconn *con) { - PGresult *res; - /* - * get the scaling factor that should be same as count(*) from - * pgbench_branches if this is not a custom query + * Note: TPC-B requires at least 100 bytes per row, and the "filler" + * fields in these table declarations were intended to comply with that. + * The pgbench_accounts table complies with that because the "filler" + * column is set to blank-padded empty string. But for all other tables + * the columns default to NULL and so don't actually take any space. We + * could fix that by giving them non-null default values. However, that + * would completely break comparability of pgbench results with prior + * versions. Since pgbench has never pretended to be fully TPC-B compliant + * anyway, we stick with the historical behavior. */ - res = PQexec(con, "select count(*) from pgbench_branches"); - if (PQresultStatus(res) != PGRES_TUPLES_OK) + struct ddlinfo { - char *sqlState = PQresultErrorField(res, PG_DIAG_SQLSTATE); - - pg_log_error("could not count number of branches: %s", PQerrorMessage(con)); + const char *table; /* table name */ + const char *smcols; /* column decls if accountIDs are 32 bits */ + const char *bigcols; /* column decls if accountIDs are 64 bits */ + int declare_fillfactor; + }; + static const struct ddlinfo DDLs[] = { + { + "pgbench_history", + "tid int,bid int,aid int,delta int,mtime timestamp,filler char(22)", + "tid int,bid int,aid bigint,delta int,mtime timestamp,filler char(22)", + 0 + }, + { + "pgbench_tellers", + "tid int not null,bid int,tbalance int,filler char(84)", + "tid int not null,bid int,tbalance int,filler char(84)", + 1 + }, + { + "pgbench_accounts", + "aid int not null,bid int,abalance int,filler char(84)", + "aid bigint not null,bid int,abalance int,filler char(84)", + 1 + }, + { + "pgbench_branches", + "bid int not null,bbalance int,filler char(88)", + "bid int not null,bbalance int,filler char(88)", + 1 + } + }; + PQExpBufferData query; - if (sqlState && strcmp(sqlState, ERRCODE_UNDEFINED_TABLE) == 0) - pg_log_error_hint("Perhaps you need to do initialization (\"pgbench -i\") in database \"%s\".", - PQdb(con)); + fprintf(stderr, "creating tables...\n"); - exit(1); - } - scale = atoi(PQgetvalue(res, 0, 0)); - if (scale < 0) - pg_fatal("invalid count(*) from pgbench_branches: \"%s\"", - PQgetvalue(res, 0, 0)); - PQclear(res); + initPQExpBuffer(&query); - /* warn if we override user-given -s switch */ - if (scale_given) - pg_log_warning("scale option ignored, using count from pgbench_branches table (%d)", - scale); + for (size_t i = 0; i < lengthof(DDLs); i++) + { + const struct ddlinfo *ddl = &DDLs[i]; - /* - * Get the partition information for the first "pgbench_accounts" table - * found in search_path. - * - * The result is empty if no "pgbench_accounts" is found. - * - * Otherwise, it always returns one row even if the table is not - * partitioned (in which case the partition strategy is NULL). - * - * The number of partitions can be 0 even for partitioned tables, if no - * partition is attached. - * - * We assume no partitioning on any failure, so as to avoid failing on an - * old version without "pg_partitioned_table". - */ - res = PQexec(con, - "select o.n, p.partstrat, pg_catalog.count(i.inhparent) " - "from pg_catalog.pg_class as c " - "join pg_catalog.pg_namespace as n on (n.oid = c.relnamespace) " - "cross join lateral (select pg_catalog.array_position(pg_catalog.current_schemas(true), n.nspname)) as o(n) " - "left join pg_catalog.pg_partitioned_table as p on (p.partrelid = c.oid) " - "left join pg_catalog.pg_inherits as i on (c.oid = i.inhparent) " - "where c.relname = 'pgbench_accounts' and o.n is not null " - "group by 1, 2 " - "order by 1 asc " - "limit 1"); + /* Construct new create table statement. */ + printfPQExpBuffer(&query, "create%s table %s(%s)", + (unlogged_tables && partition_method == PART_NONE) ? " unlogged" : "", + ddl->table, + (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols); - if (PQresultStatus(res) != PGRES_TUPLES_OK) - { - /* probably an older version, coldly assume no partitioning */ - partition_method = PART_NONE; - partitions = 0; - } - else if (PQntuples(res) == 0) - { - /* - * This case is unlikely as pgbench already found "pgbench_branches" - * above to compute the scale. - */ - pg_log_error("no pgbench_accounts table found in \"search_path\""); - pg_log_error_hint("Perhaps you need to do initialization (\"pgbench -i\") in database \"%s\".", PQdb(con)); - exit(1); - } - else /* PQntuples(res) == 1 */ - { - /* normal case, extract partition information */ - if (PQgetisnull(res, 0, 1)) - partition_method = PART_NONE; - else + /* Partition pgbench_accounts table */ + if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0) + appendPQExpBuffer(&query, + " partition by %s (aid)", PARTITION_METHOD[partition_method]); + else if (ddl->declare_fillfactor) { - char *ps = PQgetvalue(res, 0, 1); + /* fillfactor is only expected on actual tables */ + appendPQExpBuffer(&query, " with (fillfactor=%d)", fillfactor); + } - /* column must be there */ - Assert(ps != NULL); + if (tablespace != NULL) + { + char *escape_tablespace; - if (strcmp(ps, "r") == 0) - partition_method = PART_RANGE; - else if (strcmp(ps, "h") == 0) - partition_method = PART_HASH; - else - { - /* possibly a newer version with new partition method */ - pg_fatal("unexpected partition method: \"%s\"", ps); - } + escape_tablespace = PQescapeIdentifier(con, tablespace, strlen(tablespace)); + appendPQExpBuffer(&query, " tablespace %s", escape_tablespace); + PQfreemem(escape_tablespace); } - partitions = atoi(PQgetvalue(res, 0, 2)); + executeStatement(con, query.data); } - PQclear(res); + termPQExpBuffer(&query); + + if (partition_method != PART_NONE) + createPartitions(con); } /* - * Replace :param with $n throughout the command's SQL text, which - * is a modifiable string in cmd->lines. + * Truncate away any old data, in one command in case there are foreign keys */ -static bool -parseQuery(Command *cmd) +static void +initTruncateTables(PGconn *con) +{ + executeStatement(con, "truncate table " + "pgbench_accounts, " + "pgbench_branches, " + "pgbench_history, " + "pgbench_tellers"); +} + +static void +initBranch(PQExpBufferData *sql, int64 curr) +{ + /* "filler" column uses NULL */ + printfPQExpBuffer(sql, + INT64_FORMAT "\t0\t\\N\n", + curr + 1); +} + +static void +initTeller(PQExpBufferData *sql, int64 curr) { - char *sql, - *p; + /* "filler" column uses NULL */ + printfPQExpBuffer(sql, + INT64_FORMAT "\t" INT64_FORMAT "\t0\t\\N\n", + curr + 1, curr / ntellers + 1); +} - cmd->argc = 1; +static void +initAccount(PQExpBufferData *sql, int64 curr) +{ + /* "filler" column defaults to blank padded empty string */ + printfPQExpBuffer(sql, + INT64_FORMAT "\t" INT64_FORMAT "\t0\t\n", + curr + 1, curr / naccounts + 1); +} - p = sql = pg_strdup(cmd->lines.data); - while ((p = strchr(p, ':')) != NULL) - { - char var[13]; - char *name; - int eaten; +static void +initPopulateTable(PGconn *con, const char *table, int64 base, + initRowMethod init_row) +{ + int n; + int64 k; + int chars = 0; + int prev_chars = 0; + PGresult *res; + PQExpBufferData sql; + char copy_statement[256]; + const char *copy_statement_fmt = "copy %s from stdin"; + int64 total = base * scale; - name = parseVariable(p, &eaten); - if (name == NULL) - { - while (*p == ':') - { - p++; - } - continue; - } + /* used to track elapsed time and estimate of the remaining time */ + pg_time_usec_t start; + int log_interval = 1; - /* - * cmd->argv[0] is the SQL statement itself, so the max number of - * arguments is one less than MAX_ARGS - */ - if (cmd->argc >= MAX_ARGS) - { - pg_log_error("statement has too many arguments (maximum is %d): %s", - MAX_ARGS - 1, cmd->lines.data); - pg_free(name); - return false; - } + /* Stay on the same line if reporting to a terminal */ + char eol = isatty(fileno(stderr)) ? '\r' : '\n'; - sprintf(var, "$%d", cmd->argc); - p = replaceVariable(&sql, p, eaten, var); + initPQExpBuffer(&sql); - cmd->argv[cmd->argc] = name; - cmd->argc++; - } + /* Use COPY with FREEZE on v14 and later for all ordinary tables */ + if ((PQserverVersion(con) >= 140000) && + get_table_relkind(con, table) == RELKIND_RELATION) + copy_statement_fmt = "copy %s from stdin with (freeze on)"; + + + n = pg_snprintf(copy_statement, sizeof(copy_statement), copy_statement_fmt, table); + if (n >= sizeof(copy_statement)) + pg_fatal("invalid buffer size: must be at least %d characters long", n); + else if (n == -1) + pg_fatal("invalid format string"); + + res = PQexec(con, copy_statement); + + if (PQresultStatus(res) != PGRES_COPY_IN) + pg_fatal("unexpected copy in result: %s", PQerrorMessage(con)); + PQclear(res); - Assert(cmd->argv[0] == NULL); - cmd->argv[0] = sql; - return true; -} + start = pg_time_now(); -/* - * syntax error while parsing a script (in practice, while parsing a - * backslash command, because we don't detect syntax errors in SQL) - * - * source: source of script (filename or builtin-script ID) - * lineno: line number within script (count from 1) - * line: whole line of backslash command, if available - * command: backslash command name, if available - * msg: the actual error message - * more: optional extra message - * column: zero-based column number, or -1 if unknown - */ -void -syntax_error(const char *source, int lineno, - const char *line, const char *command, - const char *msg, const char *more, int column) -{ - PQExpBufferData buf; + for (k = 0; k < total; k++) + { + int64 j = k + 1; - initPQExpBuffer(&buf); + init_row(&sql, k); + if (PQputline(con, sql.data)) + pg_fatal("PQputline failed"); - printfPQExpBuffer(&buf, "%s:%d: %s", source, lineno, msg); - if (more != NULL) - appendPQExpBuffer(&buf, " (%s)", more); - if (column >= 0 && line == NULL) - appendPQExpBuffer(&buf, " at column %d", column + 1); - if (command != NULL) - appendPQExpBuffer(&buf, " in command \"%s\"", command); + if (CancelRequested) + break; - pg_log_error("%s", buf.data); + /* + * If we want to stick with the original logging, print a message each + * 100k inserted rows. + */ + if ((!use_quiet) && (j % 100000 == 0)) + { + double elapsed_sec = PG_TIME_GET_DOUBLE(pg_time_now() - start); + double remaining_sec = ((double) total - j) * elapsed_sec / j; - termPQExpBuffer(&buf); + chars = fprintf(stderr, INT64_FORMAT " of " INT64_FORMAT " tuples (%d%%) of %s done (elapsed %.2f s, remaining %.2f s)", + j, total, + (int) ((j * 100) / total), + table, elapsed_sec, remaining_sec); - if (line != NULL) - { - fprintf(stderr, "%s\n", line); - if (column >= 0) - fprintf(stderr, "%*c error found here\n", column + 1, '^'); - } + /* + * If the previous progress message is longer than the current + * one, add spaces to the current line to fully overwrite any + * remaining characters from the previous message. + */ + if (prev_chars > chars) + fprintf(stderr, "%*c", prev_chars - chars, ' '); + fputc(eol, stderr); + prev_chars = chars; + } + /* let's not call the timing for each row, but only each 100 rows */ + else if (use_quiet && (j % 100 == 0)) + { + double elapsed_sec = PG_TIME_GET_DOUBLE(pg_time_now() - start); + double remaining_sec = ((double) total - j) * elapsed_sec / j; - exit(1); -} + /* have we reached the next interval (or end)? */ + if ((j == total) || (elapsed_sec >= log_interval * LOG_STEP_SECONDS)) + { + chars = fprintf(stderr, INT64_FORMAT " of " INT64_FORMAT " tuples (%d%%) of %s done (elapsed %.2f s, remaining %.2f s)", + j, total, + (int) ((j * 100) / total), + table, elapsed_sec, remaining_sec); -/* - * Return a pointer to the start of the SQL command, after skipping over - * whitespace and "--" comments. - * If the end of the string is reached, return NULL. - */ -static char * -skip_sql_comments(char *sql_command) -{ - char *p = sql_command; + /* + * If the previous progress message is longer than the current + * one, add spaces to the current line to fully overwrite any + * remaining characters from the previous message. + */ + if (prev_chars > chars) + fprintf(stderr, "%*c", prev_chars - chars, ' '); + fputc(eol, stderr); + prev_chars = chars; - /* Skip any leading whitespace, as well as "--" style comments */ - for (;;) - { - if (isspace((unsigned char) *p)) - p++; - else if (strncmp(p, "--", 2) == 0) - { - p = strchr(p, '\n'); - if (p == NULL) - return NULL; - p++; + /* skip to the next interval */ + log_interval = (int) ceil(elapsed_sec / LOG_STEP_SECONDS); + } } - else - break; } - /* NULL if there's nothing but whitespace and comments */ - if (*p == '\0') - return NULL; + if (chars != 0 && eol != '\n') + fprintf(stderr, "%*c\r", chars, ' '); /* Clear the current line */ + + if (PQputline(con, "\\.\n")) + pg_fatal("very last PQputline failed"); + if (PQendcopy(con)) + pg_fatal("PQendcopy failed"); - return p; + termPQExpBuffer(&sql); } /* - * Parse a SQL command; return a Command struct, or NULL if it's a comment + * Fill the standard tables with some data generated and sent from the client. * - * On entry, psqlscan.l has collected the command into "buf", so we don't - * really need to do much here except check for comments and set up a Command - * struct. + * The filler column is NULL in pgbench_branches and pgbench_tellers, and is + * a blank-padded string in pgbench_accounts. */ -static Command * -create_sql_command(PQExpBuffer buf) -{ - Command *my_command; - char *p = skip_sql_comments(buf->data); - - if (p == NULL) - return NULL; - - /* Allocate and initialize Command structure */ - my_command = pg_malloc0_object(Command); - initPQExpBuffer(&my_command->lines); - appendPQExpBufferStr(&my_command->lines, p); - my_command->first_line = NULL; /* this is set later */ - my_command->type = SQL_COMMAND; - my_command->meta = META_NONE; - my_command->argc = 0; - my_command->retries = 0; - my_command->failures = 0; - memset(my_command->argv, 0, sizeof(my_command->argv)); - my_command->varprefix = NULL; /* allocated later, if needed */ - my_command->expr = NULL; - initSimpleStats(&my_command->stats); - my_command->prepname = NULL; /* set later, if needed */ - - return my_command; -} - -/* Free a Command structure and associated data */ static void -free_command(Command *command) +initGenerateDataClientSide(PGconn *con) { - termPQExpBuffer(&command->lines); - pg_free(command->first_line); - for (int i = 0; i < command->argc; i++) - pg_free(command->argv[i]); - pg_free(command->varprefix); + fprintf(stderr, "generating data (client-side)...\n"); /* - * It should also free expr recursively, but this is currently not needed - * as only gset commands (which do not have an expression) are freed. + * we do all of this in one transaction to enable the backend's + * data-loading optimizations */ - pg_free(command); -} - -/* - * Once an SQL command is fully parsed, possibly by accumulating several - * parts, complete other fields of the Command structure. - */ -static void -postprocess_sql_command(Command *my_command) -{ - char buffer[128]; - static int prepnum = 0; + executeStatement(con, "begin"); - Assert(my_command->type == SQL_COMMAND); + /* truncate away any old data */ + initTruncateTables(con); - /* Save the first line for error display. */ - strlcpy(buffer, my_command->lines.data, sizeof(buffer)); - buffer[strcspn(buffer, "\n\r")] = '\0'; - my_command->first_line = pg_strdup(buffer); + /* + * fill branches, tellers, accounts in that order in case foreign keys + * already exist + */ + initPopulateTable(con, "pgbench_branches", nbranches, initBranch); + initPopulateTable(con, "pgbench_tellers", ntellers, initTeller); + initPopulateTable(con, "pgbench_accounts", naccounts, initAccount); - /* Parse query and generate prepared statement name, if necessary */ - switch (querymode) - { - case QUERY_SIMPLE: - my_command->argv[0] = my_command->lines.data; - my_command->argc++; - break; - case QUERY_PREPARED: - my_command->prepname = psprintf("P_%d", prepnum++); - pg_fallthrough; - case QUERY_EXTENDED: - if (!parseQuery(my_command)) - exit(1); - break; - default: - exit(1); - } + executeStatement(con, "commit"); } /* - * Parse a backslash command; return a Command struct, or NULL if comment + * Fill the standard tables with some data generated on the server * - * At call, we have scanned only the initial backslash. + * As already the case with the client-side data generation, the filler + * column defaults to NULL in pgbench_branches and pgbench_tellers, + * and is a blank-padded string in pgbench_accounts. */ -static Command * -process_backslash_command(PsqlScanState sstate, const char *source, - int lineno, int start_offset) +static void +initGenerateDataServerSide(PGconn *con) { - Command *my_command; - PQExpBufferData word_buf; - int word_offset; - int offsets[MAX_ARGS]; /* offsets of argument words */ - int j; - - initPQExpBuffer(&word_buf); - - /* Collect first word of command */ - if (!expr_lex_one_word(sstate, &word_buf, &word_offset)) - { - termPQExpBuffer(&word_buf); - return NULL; - } - - /* Allocate and initialize Command structure */ - my_command = pg_malloc0_object(Command); - my_command->type = META_COMMAND; - my_command->argc = 0; - initSimpleStats(&my_command->stats); - - /* Save first word (command name) */ - j = 0; - offsets[j] = word_offset; - my_command->argv[j++] = pg_strdup(word_buf.data); - my_command->argc++; - - /* ... and convert it to enum form */ - my_command->meta = getMetaCommand(my_command->argv[0]); + PQExpBufferData sql; - if (my_command->meta == META_SET || - my_command->meta == META_IF || - my_command->meta == META_ELIF) - { - yyscan_t yyscanner; + fprintf(stderr, "generating data (server-side)...\n"); - /* For \set, collect var name */ - if (my_command->meta == META_SET) - { - if (!expr_lex_one_word(sstate, &word_buf, &word_offset)) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "missing argument", NULL, -1); + /* + * we do all of this in one transaction to enable the backend's + * data-loading optimizations + */ + executeStatement(con, "begin"); - offsets[j] = word_offset; - my_command->argv[j++] = pg_strdup(word_buf.data); - my_command->argc++; - } + /* truncate away any old data */ + initTruncateTables(con); - /* then for all parse the expression */ - yyscanner = expr_scanner_init(sstate, source, lineno, start_offset, - my_command->argv[0]); + initPQExpBuffer(&sql); - if (expr_yyparse(&my_command->expr, yyscanner) != 0) - { - /* dead code: exit done from syntax_error called by yyerror */ - exit(1); - } + printfPQExpBuffer(&sql, + "insert into pgbench_branches(bid,bbalance) " + "select bid, 0 " + "from generate_series(1, %d) as bid", nbranches * scale); + executeStatement(con, sql.data); - /* Save line, trimming any trailing newline */ - my_command->first_line = - expr_scanner_get_substring(sstate, - start_offset, - true); + printfPQExpBuffer(&sql, + "insert into pgbench_tellers(tid,bid,tbalance) " + "select tid, (tid - 1) / %d + 1, 0 " + "from generate_series(1, %d) as tid", ntellers, ntellers * scale); + executeStatement(con, sql.data); - expr_scanner_finish(yyscanner); + printfPQExpBuffer(&sql, + "insert into pgbench_accounts(aid,bid,abalance,filler) " + "select aid, (aid - 1) / %d + 1, 0, '' " + "from generate_series(1, " INT64_FORMAT ") as aid", + naccounts, (int64) naccounts * scale); + executeStatement(con, sql.data); - termPQExpBuffer(&word_buf); + termPQExpBuffer(&sql); - return my_command; - } + executeStatement(con, "commit"); +} - /* For all other commands, collect remaining words. */ - while (expr_lex_one_word(sstate, &word_buf, &word_offset)) - { - /* - * my_command->argv[0] is the command itself, so the max number of - * arguments is one less than MAX_ARGS - */ - if (j >= MAX_ARGS) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "too many arguments", NULL, -1); +/* + * Invoke vacuum on the standard tables + */ +static void +initVacuum(PGconn *con) +{ + fprintf(stderr, "vacuuming...\n"); + executeStatement(con, "vacuum analyze pgbench_branches"); + executeStatement(con, "vacuum analyze pgbench_tellers"); + executeStatement(con, "vacuum analyze pgbench_accounts"); + executeStatement(con, "vacuum analyze pgbench_history"); +} - offsets[j] = word_offset; - my_command->argv[j++] = pg_strdup(word_buf.data); - my_command->argc++; - } +/* + * Create primary keys on the standard tables + */ +static void +initCreatePKeys(PGconn *con) +{ + static const char *const DDLINDEXes[] = { + "alter table pgbench_branches add primary key (bid)", + "alter table pgbench_tellers add primary key (tid)", + "alter table pgbench_accounts add primary key (aid)" + }; + PQExpBufferData query; - /* Save line, trimming any trailing newline */ - my_command->first_line = - expr_scanner_get_substring(sstate, - start_offset, - true); + fprintf(stderr, "creating primary keys...\n"); + initPQExpBuffer(&query); - if (my_command->meta == META_SLEEP) + for (size_t i = 0; i < lengthof(DDLINDEXes); i++) { - if (my_command->argc < 2) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "missing argument", NULL, -1); - - if (my_command->argc > 3) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "too many arguments", NULL, - offsets[3] - start_offset); + resetPQExpBuffer(&query); + appendPQExpBufferStr(&query, DDLINDEXes[i]); - /* - * Split argument into number and unit to allow "sleep 1ms" etc. We - * don't have to terminate the number argument with null because it - * will be parsed with atoi, which ignores trailing non-digit - * characters. - */ - if (my_command->argv[1][0] != ':') + if (index_tablespace != NULL) { - 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; - - /* Eat all digits */ - while (*c && isdigit((unsigned char) *c)) - c++; + char *escape_tablespace; - if (*c) - { - 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); - } - } + escape_tablespace = PQescapeIdentifier(con, index_tablespace, + strlen(index_tablespace)); + appendPQExpBuffer(&query, " using index tablespace %s", escape_tablespace); + PQfreemem(escape_tablespace); } - if (my_command->argc == 3) - { - if (pg_strcasecmp(my_command->argv[2], "us") != 0 && - pg_strcasecmp(my_command->argv[2], "ms") != 0 && - pg_strcasecmp(my_command->argv[2], "s") != 0) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "unrecognized time unit, must be us, ms or s", - my_command->argv[2], offsets[2] - start_offset); - } - } - else if (my_command->meta == META_SETSHELL) - { - if (my_command->argc < 3) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "missing argument", NULL, -1); - } - else if (my_command->meta == META_SHELL) - { - if (my_command->argc < 2) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "missing command", NULL, -1); - } - else if (my_command->meta == META_ELSE || my_command->meta == META_ENDIF || - my_command->meta == META_STARTPIPELINE || - my_command->meta == META_ENDPIPELINE || - my_command->meta == META_SYNCPIPELINE) - { - if (my_command->argc != 1) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "unexpected argument", NULL, -1); - } - else if (my_command->meta == META_GSET || my_command->meta == META_ASET) - { - if (my_command->argc > 2) - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "too many arguments", NULL, -1); - } - else - { - /* my_command->meta == META_NONE */ - syntax_error(source, lineno, my_command->first_line, my_command->argv[0], - "invalid command", NULL, -1); + executeStatement(con, query.data); } - termPQExpBuffer(&word_buf); - - return my_command; + termPQExpBuffer(&query); } +/* + * Create foreign key constraints between the standard tables + */ static void -ConditionError(const char *desc, int cmdn, const char *msg) +initCreateFKeys(PGconn *con) { - pg_fatal("condition error in script \"%s\" command %d: %s", - desc, cmdn, msg); + static const char *const DDLKEYs[] = { + "alter table pgbench_tellers add constraint pgbench_tellers_bid_fkey foreign key (bid) references pgbench_branches", + "alter table pgbench_accounts add constraint pgbench_accounts_bid_fkey foreign key (bid) references pgbench_branches", + "alter table pgbench_history add constraint pgbench_history_bid_fkey foreign key (bid) references pgbench_branches", + "alter table pgbench_history add constraint pgbench_history_tid_fkey foreign key (tid) references pgbench_tellers", + "alter table pgbench_history add constraint pgbench_history_aid_fkey foreign key (aid) references pgbench_accounts" + }; + + fprintf(stderr, "creating foreign keys...\n"); + for (size_t i = 0; i < lengthof(DDLKEYs); i++) + { + executeStatement(con, DDLKEYs[i]); + } } /* - * Partial evaluation of conditionals before recording and running the script. + * Validate an initialization-steps string + * + * (We could just leave it to runInitSteps() to fail if there are wrong + * characters, but since initialization can take awhile, it seems friendlier + * to check during option parsing.) */ static void -CheckConditional(const ParsedScript *ps) +checkInitSteps(const char *initialize_steps) { - /* statically check conditional structure */ - ConditionalStack cs = conditional_stack_create(); - int i; + if (initialize_steps[0] == '\0') + pg_fatal("no initialization steps specified"); - for (i = 0; ps->commands[i] != NULL; i++) + for (const char *step = initialize_steps; *step != '\0'; step++) { - Command *cmd = ps->commands[i]; - - if (cmd->type == META_COMMAND) + if (strchr(ALL_INIT_STEPS " ", *step) == NULL) { - switch (cmd->meta) - { - case META_IF: - conditional_stack_push(cs, IFSTATE_FALSE); - break; - case META_ELIF: - if (conditional_stack_empty(cs)) - ConditionError(ps->desc, i + 1, "\\elif without matching \\if"); - if (conditional_stack_peek(cs) == IFSTATE_ELSE_FALSE) - ConditionError(ps->desc, i + 1, "\\elif after \\else"); - break; - case META_ELSE: - if (conditional_stack_empty(cs)) - ConditionError(ps->desc, i + 1, "\\else without matching \\if"); - if (conditional_stack_peek(cs) == IFSTATE_ELSE_FALSE) - ConditionError(ps->desc, i + 1, "\\else after \\else"); - conditional_stack_poke(cs, IFSTATE_ELSE_FALSE); - break; - case META_ENDIF: - if (!conditional_stack_pop(cs)) - ConditionError(ps->desc, i + 1, "\\endif without matching \\if"); - break; - default: - /* ignore anything else... */ - break; - } + pg_log_error("unrecognized initialization step \"%c\"", *step); + pg_log_error_detail("Allowed step characters are: \"" ALL_INIT_STEPS "\"."); + exit(1); } } - if (!conditional_stack_empty(cs)) - ConditionError(ps->desc, i + 1, "\\if without matching \\endif"); - conditional_stack_destroy(cs); } /* - * Parse a script (either the contents of a file, or a built-in script) - * and add it to the list of scripts. + * Invoke each initialization step in the given string */ static void -ParseScript(const char *script, const char *desc, int weight) +runInitSteps(const char *initialize_steps) { - ParsedScript ps; - PsqlScanState sstate; - PQExpBufferData line_buf; - int alloc_num; - int index; - -#define COMMANDS_ALLOC_NUM 128 - alloc_num = COMMANDS_ALLOC_NUM; - - /* Initialize all fields of ps */ - ps.desc = desc; - ps.weight = weight; - ps.commands = pg_malloc_array(Command *, alloc_num); - initStats(&ps.stats, 0); - - /* Prepare to parse script */ - sstate = psql_scan_create(&pgbench_callbacks); + PQExpBufferData stats; + PGconn *con; + const char *step; + double run_time = 0.0; + bool first = true; - /* - * Ideally, we'd scan scripts using the encoding and stdstrings settings - * we get from a DB connection. However, without major rearrangement of - * pgbench's argument parsing, we can't have a DB connection at the time - * we parse scripts. Using SQL_ASCII (encoding 0) should work well enough - * with any backend-safe encoding, though conceivably we could be fooled - * if a script file uses a client-only encoding. We also assume that - * stdstrings should be true, which is a bit riskier. - */ - psql_scan_setup(sstate, script, strlen(script), 0, true); + initPQExpBuffer(&stats); - initPQExpBuffer(&line_buf); + if ((con = doConnect()) == NULL) + pg_fatal("could not create connection for initialization"); - index = 0; + setup_cancel_handler(NULL); + SetCancelConn(con); - for (;;) + for (step = initialize_steps; *step != '\0'; step++) { - PsqlScanResult sr; - promptStatus_t prompt; - Command *command = NULL; - - resetPQExpBuffer(&line_buf); - - sr = psql_scan(sstate, &line_buf, &prompt); - - /* If we collected a new SQL command, process that */ - command = create_sql_command(&line_buf); - - /* store new command */ - if (command) - ps.commands[index++] = command; + char *op = NULL; + pg_time_usec_t start = pg_time_now(); - /* If we reached a backslash, process that */ - if (sr == PSCAN_BACKSLASH) + switch (*step) { - int lineno; - int start_offset; - - /* Capture location of the backslash */ - psql_scan_get_location(sstate, &lineno, &start_offset); - start_offset--; - - command = process_backslash_command(sstate, desc, - lineno, start_offset); - - if (command) - { - /* - * If this is gset or aset, merge into the preceding command. - * (We don't use a command slot in this case). - */ - if (command->meta == META_GSET || command->meta == META_ASET) - { - Command *cmd; - - if (index == 0) - syntax_error(desc, lineno, NULL, NULL, - "\\gset must follow an SQL command", - NULL, -1); - - cmd = ps.commands[index - 1]; - - if (cmd->type != SQL_COMMAND || - cmd->varprefix != NULL) - syntax_error(desc, lineno, NULL, NULL, - "\\gset must follow an SQL command", - cmd->first_line, -1); - - /* get variable prefix */ - if (command->argc <= 1 || command->argv[1][0] == '\0') - cmd->varprefix = pg_strdup(""); - else - cmd->varprefix = pg_strdup(command->argv[1]); - - /* update the sql command meta */ - cmd->meta = command->meta; - - /* cleanup unused command */ - free_command(command); - - continue; - } - - /* Attach any other backslash command as a new command */ - ps.commands[index++] = command; - } + case 'd': + op = "drop tables"; + initDropTables(con); + break; + case 't': + op = "create tables"; + initCreateTables(con); + break; + case 'g': + op = "client-side generate"; + initGenerateDataClientSide(con); + break; + case 'G': + op = "server-side generate"; + initGenerateDataServerSide(con); + break; + case 'v': + op = "vacuum"; + initVacuum(con); + break; + case 'p': + op = "primary keys"; + initCreatePKeys(con); + break; + case 'f': + op = "foreign keys"; + initCreateFKeys(con); + break; + case ' ': + break; /* ignore */ + default: + pg_log_error("unrecognized initialization step \"%c\"", *step); + PQfinish(con); + exit(1); } - /* - * Since we used a command slot, allocate more if needed. Note we - * always allocate one more in order to accommodate the NULL - * terminator below. - */ - if (index >= alloc_num) + if (op != NULL) { - alloc_num += COMMANDS_ALLOC_NUM; - ps.commands = (Command **) - pg_realloc_array(ps.commands, Command *, alloc_num); - } - - /* Done if we reached EOF */ - if (sr == PSCAN_INCOMPLETE || sr == PSCAN_EOL) - break; - } - - ps.commands[index] = NULL; - - addScript(&ps); - - termPQExpBuffer(&line_buf); - psql_scan_finish(sstate); - psql_scan_destroy(sstate); -} - -/* - * Read the entire contents of file fd, and return it in a malloc'd buffer. - * - * The buffer will typically be larger than necessary, but we don't care - * in this program, because we'll free it as soon as we've parsed the script. - */ -static char * -read_file_contents(FILE *fd) -{ - char *buf; - size_t buflen = BUFSIZ; - size_t used = 0; + double elapsed_sec = PG_TIME_GET_DOUBLE(pg_time_now() - start); - buf = (char *) pg_malloc(buflen); + if (!first) + appendPQExpBufferStr(&stats, ", "); + else + first = false; - for (;;) - { - size_t nread; + appendPQExpBuffer(&stats, "%s %.2f s", op, elapsed_sec); - nread = fread(buf + used, 1, BUFSIZ, fd); - used += nread; - /* If fread() read less than requested, must be EOF or error */ - if (nread < BUFSIZ) - break; - /* Enlarge buf so we can read some more */ - buflen += BUFSIZ; - buf = (char *) pg_realloc(buf, buflen); + run_time += elapsed_sec; + } } - /* There is surely room for a terminator */ - buf[used] = '\0'; - return buf; + fprintf(stderr, "done in %.2f s (%s).\n", run_time, stats.data); + ResetCancelConn(); + PQfinish(con); + termPQExpBuffer(&stats); } /* - * Given a file name, read it and add its script to the list. - * "-" means to read stdin. - * NB: filename must be storage that won't disappear. + * Extract pgbench table information into global variables scale, + * partition_method and partitions. */ static void -process_file(const char *filename, int weight) +GetTableInfo(PGconn *con, bool scale_given) { - FILE *fd; - char *buf; - - /* Slurp the file contents into "buf" */ - if (strcmp(filename, "-") == 0) - fd = stdin; - else if ((fd = fopen(filename, "r")) == NULL) - pg_fatal("could not open file \"%s\": %m", filename); - - buf = read_file_contents(fd); - - if (ferror(fd)) - pg_fatal("could not read file \"%s\": %m", filename); - - if (fd != stdin) - fclose(fd); - - ParseScript(buf, filename, weight); - - free(buf); -} + PGresult *res; -/* Parse the given builtin script and add it to the list. */ -static void -process_builtin(const BuiltinScript *bi, int weight) -{ - ParseScript(bi->script, bi->desc, weight); -} + /* + * get the scaling factor that should be same as count(*) from + * pgbench_branches if this is not a custom query + */ + res = PQexec(con, "select count(*) from pgbench_branches"); + if (PQresultStatus(res) != PGRES_TUPLES_OK) + { + char *sqlState = PQresultErrorField(res, PG_DIAG_SQLSTATE); -/* show available builtin scripts */ -static void -listAvailableScripts(void) -{ - fprintf(stderr, "Available builtin scripts:\n"); - for (size_t i = 0; i < lengthof(builtin_script); i++) - fprintf(stderr, " %13s: %s\n", builtin_script[i].name, builtin_script[i].desc); - fprintf(stderr, "\n"); -} + pg_log_error("could not count number of branches: %s", PQerrorMessage(con)); -/* return builtin script "name" if unambiguous, fails if not found */ -static const BuiltinScript * -findBuiltin(const char *name) -{ - int found = 0, - len = strlen(name); - const BuiltinScript *result = NULL; + if (sqlState && strcmp(sqlState, ERRCODE_UNDEFINED_TABLE) == 0) + pg_log_error_hint("Perhaps you need to do initialization (\"pgbench -i\") in database \"%s\".", + PQdb(con)); - for (size_t i = 0; i < lengthof(builtin_script); i++) - { - if (strncmp(builtin_script[i].name, name, len) == 0) - { - result = &builtin_script[i]; - found++; - } + exit(1); } + scale = atoi(PQgetvalue(res, 0, 0)); + if (scale < 0) + pg_fatal("invalid count(*) from pgbench_branches: \"%s\"", + PQgetvalue(res, 0, 0)); + PQclear(res); - /* ok, unambiguous result */ - if (found == 1) - return result; - - /* error cases */ - if (found == 0) - pg_log_error("no builtin script found for name \"%s\"", name); - else /* found > 1 */ - pg_log_error("ambiguous builtin name: %d builtin scripts found for prefix \"%s\"", found, name); - - listAvailableScripts(); - exit(1); -} + /* warn if we override user-given -s switch */ + if (scale_given) + pg_log_warning("scale option ignored, using count from pgbench_branches table (%d)", + scale); -/* - * Determine the weight specification from a script option (-b, -f), if any, - * and return it as an integer (1 is returned if there's no weight). The - * script name is returned in *script as a malloc'd string. - */ -static int -parseScriptWeight(const char *option, char **script) -{ - const char *sep; - int weight; + /* + * Get the partition information for the first "pgbench_accounts" table + * found in search_path. + * + * The result is empty if no "pgbench_accounts" is found. + * + * Otherwise, it always returns one row even if the table is not + * partitioned (in which case the partition strategy is NULL). + * + * The number of partitions can be 0 even for partitioned tables, if no + * partition is attached. + * + * We assume no partitioning on any failure, so as to avoid failing on an + * old version without "pg_partitioned_table". + */ + res = PQexec(con, + "select o.n, p.partstrat, pg_catalog.count(i.inhparent) " + "from pg_catalog.pg_class as c " + "join pg_catalog.pg_namespace as n on (n.oid = c.relnamespace) " + "cross join lateral (select pg_catalog.array_position(pg_catalog.current_schemas(true), n.nspname)) as o(n) " + "left join pg_catalog.pg_partitioned_table as p on (p.partrelid = c.oid) " + "left join pg_catalog.pg_inherits as i on (c.oid = i.inhparent) " + "where c.relname = 'pgbench_accounts' and o.n is not null " + "group by 1, 2 " + "order by 1 asc " + "limit 1"); - if ((sep = strrchr(option, WSEP))) + if (PQresultStatus(res) != PGRES_TUPLES_OK) { - int namelen = sep - option; - long wtmp; - char *badp; - - /* generate the script name */ - *script = pg_malloc(namelen + 1); - strncpy(*script, option, namelen); - (*script)[namelen] = '\0'; - - /* process digits of the weight spec */ - errno = 0; - wtmp = strtol(sep + 1, &badp, 10); - if (errno != 0 || badp == sep + 1 || *badp != '\0') - pg_fatal("invalid weight specification: %s", sep); - if (wtmp > INT_MAX || wtmp < 0) - pg_fatal("weight specification out of range (0 .. %d): %ld", - INT_MAX, wtmp); - weight = wtmp; + /* probably an older version, coldly assume no partitioning */ + partition_method = PART_NONE; + partitions = 0; } - else + else if (PQntuples(res) == 0) { - *script = pg_strdup(option); - weight = 1; + /* + * This case is unlikely as pgbench already found "pgbench_branches" + * above to compute the scale. + */ + pg_log_error("no pgbench_accounts table found in \"search_path\""); + pg_log_error_hint("Perhaps you need to do initialization (\"pgbench -i\") in database \"%s\".", PQdb(con)); + exit(1); } + else /* PQntuples(res) == 1 */ + { + /* normal case, extract partition information */ + if (PQgetisnull(res, 0, 1)) + partition_method = PART_NONE; + else + { + char *ps = PQgetvalue(res, 0, 1); - return weight; -} - -/* append a script to the list of scripts to process */ -static void -addScript(const ParsedScript *script) -{ - if (script->commands == NULL || script->commands[0] == NULL) - pg_fatal("empty command list for script \"%s\"", script->desc); + /* column must be there */ + Assert(ps != NULL); - if (num_scripts >= MAX_SCRIPTS) - pg_fatal("at most %d SQL scripts are allowed", MAX_SCRIPTS); + if (strcmp(ps, "r") == 0) + partition_method = PART_RANGE; + else if (strcmp(ps, "h") == 0) + partition_method = PART_HASH; + else + { + /* possibly a newer version with new partition method */ + pg_fatal("unexpected partition method: \"%s\"", ps); + } + } - CheckConditional(script); + partitions = atoi(PQgetvalue(res, 0, 2)); + } - sql_script[num_scripts] = *script; - num_scripts++; + PQclear(res); } + /* * Print progress report. * diff --git a/src/bin/pgbench/pgbench.h b/src/bin/pgbench/pgbench.h index 50c13277741..598a13e3b61 100644 --- a/src/bin/pgbench/pgbench.h +++ b/src/bin/pgbench/pgbench.h @@ -12,7 +12,10 @@ #define PGBENCH_H #include "fe_utils/psqlscan.h" - +#include "common/pg_prng.h" +#include "fe_utils/conditional.h" +#include "libpq-fe.h" +#include "stats.h" #include "variable.h" /* @@ -116,6 +119,253 @@ struct PgBenchExprList PgBenchExprLink *tail; }; + +/* + * Multi-platform thread implementations + */ + +#ifdef WIN32 +/* Use Windows threads */ +#include +#define GETERRNO() (_dosmaperr(GetLastError()), errno) +#define THREAD_T HANDLE +#define THREAD_FUNC_RETURN_TYPE unsigned +#define THREAD_FUNC_RETURN return 0 +#define THREAD_FUNC_CC __stdcall +#define THREAD_CREATE(handle, function, arg) \ + ((*(handle) = (HANDLE) _beginthreadex(NULL, 0, (function), (arg), 0, NULL)) == 0 ? errno : 0) +#define THREAD_JOIN(handle) \ + (WaitForSingleObject(handle, INFINITE) != WAIT_OBJECT_0 ? \ + GETERRNO() : CloseHandle(handle) ? 0 : GETERRNO()) +#define THREAD_BARRIER_T SYNCHRONIZATION_BARRIER +#define THREAD_BARRIER_INIT(barrier, n) \ + (InitializeSynchronizationBarrier((barrier), (n), 0) ? 0 : GETERRNO()) +#define THREAD_BARRIER_WAIT(barrier) \ + EnterSynchronizationBarrier((barrier), \ + SYNCHRONIZATION_BARRIER_FLAGS_BLOCK_ONLY) +#define THREAD_BARRIER_DESTROY(barrier) +#else +/* Use POSIX threads */ +#include "port/pg_pthread.h" +#define THREAD_T pthread_t +#define THREAD_FUNC_RETURN_TYPE void * +#define THREAD_FUNC_RETURN return NULL +#define THREAD_FUNC_CC +#define THREAD_CREATE(handle, function, arg) \ + pthread_create((handle), NULL, (function), (arg)) +#define THREAD_JOIN(handle) \ + pthread_join((handle), NULL) +#define THREAD_BARRIER_T pthread_barrier_t +#define THREAD_BARRIER_INIT(barrier, n) \ + pthread_barrier_init((barrier), NULL, (n)) +#define THREAD_BARRIER_WAIT(barrier) pthread_barrier_wait((barrier)) +#define THREAD_BARRIER_DESTROY(barrier) pthread_barrier_destroy((barrier)) +#endif + +/* + * Transaction status at the end of a command. + */ +typedef enum TStatus +{ + TSTATUS_IDLE, + TSTATUS_IN_BLOCK, + TSTATUS_CONN_ERROR, + TSTATUS_OTHER_ERROR, +} TStatus; + +/* + * Connection state machine states. + */ +typedef enum +{ + /* + * The client must first choose a script to execute. Once chosen, it can + * either be throttled (state CSTATE_PREPARE_THROTTLE under --rate), start + * right away (state CSTATE_START_TX) or not start at all if the timer was + * exceeded (state CSTATE_FINISHED). + */ + CSTATE_CHOOSE_SCRIPT, + + /* + * CSTATE_START_TX performs start-of-transaction processing. Establishes + * a new connection for the transaction in --connect mode, records the + * transaction start time, and proceed to the first command. + * + * Note: once a script is started, it will either error or run till its + * end, where it may be interrupted. It is not interrupted while running, + * so pgbench --time is to be understood as tx are allowed to start in + * that time, and will finish when their work is completed. + */ + CSTATE_START_TX, + + /* + * In CSTATE_PREPARE_THROTTLE state, we calculate when to begin the next + * transaction, and advance to CSTATE_THROTTLE. CSTATE_THROTTLE state + * sleeps until that moment, then advances to CSTATE_START_TX, or + * CSTATE_FINISHED if the next transaction would start beyond the end of + * the run. + */ + CSTATE_PREPARE_THROTTLE, + CSTATE_THROTTLE, + + /* + * We loop through these states, to process each command in the script: + * + * CSTATE_START_COMMAND starts the execution of a command. On a SQL + * command, the command is sent to the server, and we move to + * CSTATE_WAIT_RESULT state unless in pipeline mode. On a \sleep + * meta-command, the timer is set, and we enter the CSTATE_SLEEP state to + * wait for it to expire. Other meta-commands are executed immediately. If + * the command about to start is actually beyond the end of the script, + * advance to CSTATE_END_TX. + * + * CSTATE_WAIT_RESULT waits until we get a result set back from the server + * for the current command. + * + * CSTATE_SLEEP waits until the end of \sleep. + * + * CSTATE_END_COMMAND records the end-of-command timestamp, increments the + * command counter, and loops back to CSTATE_START_COMMAND state. + * + * CSTATE_SKIP_COMMAND is used by conditional branches which are not + * executed. It quickly skip commands that do not need any evaluation. + * This state can move forward several commands, till there is something + * to do or the end of the script. + */ + CSTATE_START_COMMAND, + CSTATE_WAIT_RESULT, + CSTATE_SLEEP, + CSTATE_END_COMMAND, + CSTATE_SKIP_COMMAND, + + /* + * States for failed commands. + * + * If the SQL/meta command fails, in CSTATE_ERROR clean up after an error: + * (1) clear the conditional stack; (2) if we have an unterminated + * (possibly failed) transaction block, send the rollback command to the + * server and wait for the result in CSTATE_WAIT_ROLLBACK_RESULT. If + * something goes wrong with rolling back, go to CSTATE_ABORTED. + * + * But if everything is ok we are ready for future transactions: if this + * is a serialization or deadlock error and we can re-execute the + * transaction from the very beginning, go to CSTATE_RETRY; otherwise go + * to CSTATE_FAILURE. + * + * In CSTATE_RETRY report an error, set the same parameters for the + * transaction execution as in the previous tries and process the first + * transaction command in CSTATE_START_COMMAND. + * + * In CSTATE_FAILURE report a failure, set the parameters for the + * transaction execution as they were before the first run of this + * transaction (except for a random state) and go to CSTATE_END_TX to + * complete this transaction. + */ + CSTATE_ERROR, + CSTATE_WAIT_ROLLBACK_RESULT, + CSTATE_RETRY, + CSTATE_FAILURE, + + /* + * CSTATE_END_TX performs end-of-transaction processing. It calculates + * latency, and logs the transaction. In --connect mode, it closes the + * current connection. + * + * Then either starts over in CSTATE_CHOOSE_SCRIPT, or enters + * CSTATE_FINISHED if we have no more work to do. + */ + CSTATE_END_TX, + + /* + * Final states. CSTATE_ABORTED means that the script execution was + * aborted because a command failed, CSTATE_FINISHED means success. + */ + CSTATE_ABORTED, + CSTATE_FINISHED, +} ConnectionStateEnum; + +/* + * Connection state. + */ +typedef struct +{ + PGconn *con; /* connection handle to DB */ + int id; /* client No. */ + ConnectionStateEnum state; /* state machine's current state. */ + ConditionalStack cstack; /* enclosing conditionals state */ + + /* + * Separate randomness for each client. This is used for random functions + * PGBENCH_RANDOM_* during the execution of the script. + */ + pg_prng_state cs_func_rs; + + int use_file; /* index in sql_script for this client */ + int command; /* command number in script */ + int num_syncs; /* number of ongoing sync commands */ + + /* client variables */ + Variables variables; + + /* various times about current transaction in microseconds */ + pg_time_usec_t txn_scheduled; /* scheduled start time of transaction */ + pg_time_usec_t sleep_until; /* scheduled start time of next cmd */ + pg_time_usec_t txn_begin; /* used for measuring schedule lag times */ + pg_time_usec_t stmt_begin; /* used for measuring statement latencies */ + + /* whether client prepared each command of each script */ + bool **prepared; + + /* + * For processing failures and repeating transactions with serialization + * or deadlock errors: + */ + EStatus estatus; /* the error status of the current transaction + * execution; this is ESTATUS_NO_ERROR if + * there were no errors */ + pg_prng_state random_state; /* random state */ + uint32 tries; /* how many times have we already tried the + * current transaction? */ + + /* per client collected stats */ + int64 cnt; /* client transaction count, for -t; skipped + * and failed transactions are also counted + * here */ +} CState; + +/* + * Thread state + */ +typedef struct +{ + int tid; /* thread id */ + THREAD_T thread; /* thread handle */ + CState *state; /* array of CState */ + int nstate; /* length of state[] */ + + /* + * Separate randomness for each thread. Each thread option uses its own + * random state to make all of them independent of each other and + * therefore deterministic at the thread level. + */ + pg_prng_state ts_choose_rs; /* random state for selecting a script */ + pg_prng_state ts_throttle_rs; /* random state for transaction throttling */ + pg_prng_state ts_sample_rs; /* random state for log sampling */ + + int64 throttle_trigger; /* previous/next throttling (us) */ + FILE *logfile; /* where to log, or NULL */ + + /* per thread collected stats in microseconds */ + pg_time_usec_t create_time; /* thread creation time */ + pg_time_usec_t started_time; /* thread is running */ + pg_time_usec_t bench_start; /* thread is benchmarking */ + pg_time_usec_t conn_duration; /* cumulated connection and disconnection + * delays */ + + StatsData stats; + int64 latency_late; /* count executed but late transactions */ +} TState; + extern int expr_yyparse(PgBenchExpr **expr_parse_result_p, yyscan_t yyscanner); extern int expr_yylex(union YYSTYPE *yylval_param, yyscan_t yyscanner); pg_noreturn extern void expr_yyerror(PgBenchExpr **expr_parse_result_p, yyscan_t yyscanner, const char *message); @@ -135,5 +385,4 @@ pg_noreturn extern void syntax_error(const char *source, int lineno, const char const char *command, const char *msg, const char *more, int column); - #endif /* PGBENCH_H */ diff --git a/src/bin/pgbench/script.c b/src/bin/pgbench/script.c new file mode 100644 index 00000000000..e4727071055 --- /dev/null +++ b/src/bin/pgbench/script.c @@ -0,0 +1,1946 @@ +/*------------------------------------------------------------------------- + * + * script.c + * Script parsing, built-in scripts, and expression evaluation for pgbench + * + * Portions Copyright (c) 1996-2026, PostgreSQL Global Development Group + * Portions Copyright (c) 1994, Regents of the University of California + * + * src/bin/pgbench/script.c + * + *------------------------------------------------------------------------- + */ +#include "postgres_fe.h" + +#include +#include +#include +#include + +#include "common/int.h" +#include "common/logging.h" +#include "common/pg_prng.h" +#include "fe_utils/conditional.h" +#include "pgbench.h" +#include "port/pg_bitutils.h" +#include "script.h" + +QueryMode querymode = QUERY_SIMPLE; +const char *const QUERYMODE[] = {"simple", "extended", "prepared"}; + +ParsedScript sql_script[MAX_SCRIPTS]; /* SQL script files */ +int num_scripts; /* number of scripts in sql_script[] */ +int64 total_weight = 0; + +static const BuiltinScript builtin_script[] = +{ + { + "tpcb-like", + "", + "\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n" + "\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n" + "\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n" + "\\set delta random(-5000, 5000)\n" + "BEGIN;\n" + "UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;\n" + "SELECT abalance FROM pgbench_accounts WHERE aid = :aid;\n" + "UPDATE pgbench_tellers SET tbalance = tbalance + :delta WHERE tid = :tid;\n" + "UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid;\n" + "INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);\n" + "END;\n" + }, + { + "simple-update", + "", + "\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n" + "\\set bid random(1, " CppAsString2(nbranches) " * :scale)\n" + "\\set tid random(1, " CppAsString2(ntellers) " * :scale)\n" + "\\set delta random(-5000, 5000)\n" + "BEGIN;\n" + "UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;\n" + "SELECT abalance FROM pgbench_accounts WHERE aid = :aid;\n" + "INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);\n" + "END;\n" + }, + { + "select-only", + "", + "\\set aid random(1, " CppAsString2(naccounts) " * :scale)\n" + "SELECT abalance FROM pgbench_accounts WHERE aid = :aid;\n" + } +}; + +static void addScript(const ParsedScript *script); + +/* callback functions for our flex lexer */ +static const PsqlScanCallbacks pgbench_callbacks = { + NULL, /* don't need get_variable functionality */ +}; + +/* + * random number generator: uniform distribution from min to max inclusive. + * + * Although the limits are expressed as int64, you can't generate the full + * int64 range in one call, because the difference of the limits mustn't + * overflow int64. This is not checked. + */ +static int64 +getrand(pg_prng_state *state, int64 min, int64 max) +{ + return min + (int64) pg_prng_uint64_range(state, 0, max - min); +} + +/* + * random number generator: exponential distribution from min to max inclusive. + * the parameter is so that the density of probability for the last cut-off max + * value is exp(-parameter). + */ +static int64 +getExponentialRand(pg_prng_state *state, int64 min, int64 max, + double parameter) +{ + double cut, + uniform, + rand; + + /* abort if wrong parameter, but must really be checked beforehand */ + Assert(parameter > 0.0); + cut = exp(-parameter); + /* pg_prng_double value in [0, 1), uniform in (0, 1] */ + uniform = 1.0 - pg_prng_double(state); + + /* + * inner expression in (cut, 1] (if parameter > 0), rand in [0, 1) + */ + Assert((1.0 - cut) != 0.0); + rand = -log(cut + (1.0 - cut) * uniform) / parameter; + /* return int64 random number within between min and max */ + return min + (int64) ((max - min + 1) * rand); +} + +/* random number generator: gaussian distribution from min to max inclusive */ +static int64 +getGaussianRand(pg_prng_state *state, int64 min, int64 max, + double parameter) +{ + double stdev; + double rand; + + /* abort if parameter is too low, but must really be checked beforehand */ + Assert(parameter >= MIN_GAUSSIAN_PARAM); + + /* + * Get normally-distributed random number in the range -parameter <= stdev + * < parameter. + * + * This loop is executed until the number is in the expected range. + * + * As the minimum parameter is 2.0, the probability of looping is low: + * sqrt(-2 ln(r)) <= 2 => r >= e^{-2} ~ 0.135, then when taking the + * average sinus multiplier as 2/pi, we have a 8.6% looping probability in + * the worst case. For a parameter value of 5.0, the looping probability + * is about e^{-5} * 2 / pi ~ 0.43%. + */ + do + { + stdev = pg_prng_double_normal(state); + } + while (stdev < -parameter || stdev >= parameter); + + /* stdev is in [-parameter, parameter), normalization to [0,1) */ + rand = (stdev + parameter) / (parameter * 2.0); + + /* return int64 random number within between min and max */ + return min + (int64) ((max - min + 1) * rand); +} + +/* + * Computing zipfian using rejection method, based on + * "Non-Uniform Random Variate Generation", + * Luc Devroye, p. 550-551, Springer 1986. + * + * This works for s > 1.0, but may perform badly for s very close to 1.0. + */ +static int64 +computeIterativeZipfian(pg_prng_state *state, int64 n, double s) +{ + double b = pow(2.0, s - 1.0); + double x, + t, + u, + v; + + /* Ensure n is sane */ + if (n <= 1) + return 1; + + while (true) + { + /* random variates */ + u = pg_prng_double(state); + v = pg_prng_double(state); + + x = floor(pow(u, -1.0 / (s - 1.0))); + + t = pow(1.0 + 1.0 / x, s - 1.0); + /* reject if too large or out of bound */ + if (v * x * (t - 1.0) / (b - 1.0) <= t / b && x <= n) + break; + } + return (int64) x; +} + +/* random number generator: zipfian distribution from min to max inclusive */ +static int64 +getZipfianRand(pg_prng_state *state, int64 min, int64 max, double s) +{ + int64 n = max - min + 1; + + /* abort if parameter is invalid */ + Assert(MIN_ZIPFIAN_PARAM <= s && s <= MAX_ZIPFIAN_PARAM); + + return min - 1 + computeIterativeZipfian(state, n, s); +} + +/* + * FNV-1a hash function + */ +static int64 +getHashFnv1a(int64 val, uint64 seed) +{ + int64 result; + int i; + + result = FNV_OFFSET_BASIS ^ seed; + for (i = 0; i < 8; ++i) + { + int32 octet = val & 0xff; + + val = val >> 8; + result = result ^ octet; + result = result * FNV_PRIME; + } + + return result; +} + +/* + * Murmur2 hash function + * + * Based on original work of Austin Appleby + * https://github.com/aappleby/smhasher/blob/master/src/MurmurHash2.cpp + */ +static int64 +getHashMurmur2(int64 val, uint64 seed) +{ + uint64 result = seed ^ MM2_MUL_TIMES_8; /* sizeof(int64) */ + uint64 k = (uint64) val; + + k *= MM2_MUL; + k ^= k >> MM2_ROT; + k *= MM2_MUL; + + result ^= k; + result *= MM2_MUL; + + result ^= result >> MM2_ROT; + result *= MM2_MUL; + result ^= result >> MM2_ROT; + + return (int64) result; +} + +/* + * Pseudorandom permutation function + * + * For small sizes, this generates each of the (size!) possible permutations + * of integers in the range [0, size) with roughly equal probability. Once + * the size is larger than 20, the number of possible permutations exceeds the + * number of distinct states of the internal pseudorandom number generator, + * and so not all possible permutations can be generated, but the permutations + * chosen should continue to give the appearance of being random. + * + * THIS FUNCTION IS NOT CRYPTOGRAPHICALLY SECURE. + * DO NOT USE FOR SUCH PURPOSE. + */ +static int64 +permute(const int64 val, const int64 isize, const int64 seed) +{ + /* using a high-end PRNG is probably overkill */ + pg_prng_state state; + uint64 size; + uint64 v; + int masklen; + uint64 mask; + int i; + + if (isize < 2) + return 0; /* nothing to permute */ + + /* Initialize prng state using the seed */ + pg_prng_seed(&state, (uint64) seed); + + /* Computations are performed on unsigned values */ + size = (uint64) isize; + v = (uint64) val % size; + + /* Mask to work modulo largest power of 2 less than or equal to size */ + masklen = pg_leftmost_one_pos64(size); + mask = (((uint64) 1) << masklen) - 1; + + /* + * Permute the input value by applying several rounds of pseudorandom + * bijective transformations. The intention here is to distribute each + * input uniformly randomly across the range, and separate adjacent inputs + * approximately uniformly randomly from each other, leading to a fairly + * random overall choice of permutation. + * + * To separate adjacent inputs, we multiply by a random number modulo + * (mask + 1), which is a power of 2. For this to be a bijection, the + * multiplier must be odd. Since this is known to lead to less randomness + * in the lower bits, we also apply a rotation that shifts the topmost bit + * into the least significant bit. In the special cases where size <= 3, + * mask = 1 and each of these operations is actually a no-op, so we also + * XOR the value with a different random number to inject additional + * randomness. Since the size is generally not a power of 2, we apply + * this bijection on overlapping upper and lower halves of the input. + * + * To distribute the inputs uniformly across the range, we then also apply + * a random offset modulo the full range. + * + * Taken together, these operations resemble a modified linear + * congruential generator, as is commonly used in pseudorandom number + * generators. The number of rounds is fairly arbitrary, but six has been + * found empirically to give a fairly good tradeoff between performance + * and uniform randomness. For small sizes it selects each of the (size!) + * possible permutations with roughly equal probability. For larger + * sizes, not all permutations can be generated, but the intended random + * spread is still produced. + */ + for (i = 0; i < 6; i++) + { + uint64 m, + r, + t; + + /* Random multiply (by an odd number), XOR and rotate of lower half */ + m = (pg_prng_uint64(&state) & mask) | 1; + r = pg_prng_uint64(&state) & mask; + if (v <= mask) + { + v = ((v * m) ^ r) & mask; + v = ((v << 1) & mask) | (v >> (masklen - 1)); + } + + /* Random multiply (by an odd number), XOR and rotate of upper half */ + m = (pg_prng_uint64(&state) & mask) | 1; + r = pg_prng_uint64(&state) & mask; + t = size - 1 - v; + if (t <= mask) + { + t = ((t * m) ^ r) & mask; + t = ((t << 1) & mask) | (t >> (masklen - 1)); + v = size - 1 - t; + } + + /* Random offset */ + r = pg_prng_uint64_range(&state, 0, size - 1); + v = (v + r) % size; + } + + return (int64) v; +} + +static bool +isLazyFunc(PgBenchFunction func) +{ + return func == PGBENCH_AND || func == PGBENCH_OR || func == PGBENCH_CASE; +} + +/* lazy evaluation of some functions */ +static bool +evalLazyFunc(CState *st, + PgBenchFunction func, PgBenchExprLink *args, PgBenchValue *retval) +{ + PgBenchValue a1, + a2; + bool ba1, + ba2; + + Assert(isLazyFunc(func) && args != NULL && args->next != NULL); + + /* args points to first condition */ + if (!evaluateExpr(st, args->expr, &a1)) + return false; + + /* second condition for AND/OR and corresponding branch for CASE */ + args = args->next; + + switch (func) + { + case PGBENCH_AND: + if (a1.type == PGBT_NULL) + { + setNullValue(retval); + return true; + } + + if (!coerceToBool(&a1, &ba1)) + return false; + + if (!ba1) + { + setBoolValue(retval, false); + return true; + } + + if (!evaluateExpr(st, args->expr, &a2)) + return false; + + if (a2.type == PGBT_NULL) + { + setNullValue(retval); + return true; + } + else if (!coerceToBool(&a2, &ba2)) + return false; + else + { + setBoolValue(retval, ba2); + return true; + } + + return true; + + case PGBENCH_OR: + + if (a1.type == PGBT_NULL) + { + setNullValue(retval); + return true; + } + + if (!coerceToBool(&a1, &ba1)) + return false; + + if (ba1) + { + setBoolValue(retval, true); + return true; + } + + if (!evaluateExpr(st, args->expr, &a2)) + return false; + + if (a2.type == PGBT_NULL) + { + setNullValue(retval); + return true; + } + else if (!coerceToBool(&a2, &ba2)) + return false; + else + { + setBoolValue(retval, ba2); + return true; + } + + case PGBENCH_CASE: + /* when true, execute branch */ + if (valueTruth(&a1)) + return evaluateExpr(st, args->expr, retval); + + /* now args contains next condition or final else expression */ + args = args->next; + + /* final else case? */ + if (args->next == NULL) + return evaluateExpr(st, args->expr, retval); + + /* no, another when, proceed */ + return evalLazyFunc(st, PGBENCH_CASE, args, retval); + + default: + /* internal error, cannot get here */ + Assert(0); + break; + } + return false; +} + +/* maximum number of function arguments */ +#define MAX_FARGS 16 + +/* + * Recursive evaluation of standard functions, + * which do not require lazy evaluation. + */ +static bool +evalStandardFunc(CState *st, + PgBenchFunction func, PgBenchExprLink *args, + PgBenchValue *retval) +{ + /* evaluate all function arguments */ + int nargs = 0; + PgBenchValue vargs[MAX_FARGS] = {0}; + PgBenchExprLink *l = args; + bool has_null = false; + + for (nargs = 0; nargs < MAX_FARGS && l != NULL; nargs++, l = l->next) + { + if (!evaluateExpr(st, l->expr, &vargs[nargs])) + return false; + has_null |= vargs[nargs].type == PGBT_NULL; + } + + if (l != NULL) + { + pg_log_error("too many function arguments, maximum is %d", MAX_FARGS); + return false; + } + + /* NULL arguments */ + if (has_null && func != PGBENCH_IS && func != PGBENCH_DEBUG) + { + setNullValue(retval); + return true; + } + + /* then evaluate function */ + switch (func) + { + /* overloaded operators */ + case PGBENCH_ADD: + case PGBENCH_SUB: + case PGBENCH_MUL: + case PGBENCH_DIV: + case PGBENCH_MOD: + case PGBENCH_EQ: + case PGBENCH_NE: + case PGBENCH_LE: + case PGBENCH_LT: + { + PgBenchValue *lval = &vargs[0], + *rval = &vargs[1]; + + Assert(nargs == 2); + + /* overloaded type management, double if some double */ + if ((lval->type == PGBT_DOUBLE || + rval->type == PGBT_DOUBLE) && func != PGBENCH_MOD) + { + double ld, + rd; + + if (!coerceToDouble(lval, &ld) || + !coerceToDouble(rval, &rd)) + return false; + + switch (func) + { + case PGBENCH_ADD: + setDoubleValue(retval, ld + rd); + return true; + + case PGBENCH_SUB: + setDoubleValue(retval, ld - rd); + return true; + + case PGBENCH_MUL: + setDoubleValue(retval, ld * rd); + return true; + + case PGBENCH_DIV: + setDoubleValue(retval, ld / rd); + return true; + + case PGBENCH_EQ: + setBoolValue(retval, ld == rd); + return true; + + case PGBENCH_NE: + setBoolValue(retval, ld != rd); + return true; + + case PGBENCH_LE: + setBoolValue(retval, ld <= rd); + return true; + + case PGBENCH_LT: + setBoolValue(retval, ld < rd); + return true; + + default: + /* cannot get here */ + Assert(0); + } + } + else /* we have integer operands, or % */ + { + int64 li, + ri, + res; + + if (!coerceToInt(lval, &li) || + !coerceToInt(rval, &ri)) + return false; + + switch (func) + { + case PGBENCH_ADD: + if (pg_add_s64_overflow(li, ri, &res)) + { + pg_log_error("bigint add out of range"); + return false; + } + setIntValue(retval, res); + return true; + + case PGBENCH_SUB: + if (pg_sub_s64_overflow(li, ri, &res)) + { + pg_log_error("bigint sub out of range"); + return false; + } + setIntValue(retval, res); + return true; + + case PGBENCH_MUL: + if (pg_mul_s64_overflow(li, ri, &res)) + { + pg_log_error("bigint mul out of range"); + return false; + } + setIntValue(retval, res); + return true; + + case PGBENCH_EQ: + setBoolValue(retval, li == ri); + return true; + + case PGBENCH_NE: + setBoolValue(retval, li != ri); + return true; + + case PGBENCH_LE: + setBoolValue(retval, li <= ri); + return true; + + case PGBENCH_LT: + setBoolValue(retval, li < ri); + return true; + + case PGBENCH_DIV: + case PGBENCH_MOD: + if (ri == 0) + { + pg_log_error("division by zero"); + return false; + } + /* special handling of -1 divisor */ + if (ri == -1) + { + if (func == PGBENCH_DIV) + { + /* overflow check (needed for INT64_MIN) */ + if (li == PG_INT64_MIN) + { + pg_log_error("bigint div out of range"); + return false; + } + else + setIntValue(retval, -li); + } + else + setIntValue(retval, 0); + return true; + } + /* else divisor is not -1 */ + if (func == PGBENCH_DIV) + setIntValue(retval, li / ri); + else /* func == PGBENCH_MOD */ + setIntValue(retval, li % ri); + + return true; + + default: + /* cannot get here */ + Assert(0); + } + } + + Assert(0); + return false; /* NOTREACHED */ + } + + /* integer bitwise operators */ + case PGBENCH_BITAND: + case PGBENCH_BITOR: + case PGBENCH_BITXOR: + case PGBENCH_LSHIFT: + case PGBENCH_RSHIFT: + { + int64 li, + ri; + + if (!coerceToInt(&vargs[0], &li) || !coerceToInt(&vargs[1], &ri)) + return false; + + if (func == PGBENCH_BITAND) + setIntValue(retval, li & ri); + else if (func == PGBENCH_BITOR) + setIntValue(retval, li | ri); + else if (func == PGBENCH_BITXOR) + setIntValue(retval, li ^ ri); + else if (func == PGBENCH_LSHIFT) + setIntValue(retval, li << ri); + else if (func == PGBENCH_RSHIFT) + setIntValue(retval, li >> ri); + else /* cannot get here */ + Assert(0); + + return true; + } + + /* logical operators */ + case PGBENCH_NOT: + { + bool b; + + if (!coerceToBool(&vargs[0], &b)) + return false; + + setBoolValue(retval, !b); + return true; + } + + /* no arguments */ + case PGBENCH_PI: + setDoubleValue(retval, M_PI); + return true; + + /* 1 overloaded argument */ + case PGBENCH_ABS: + { + PgBenchValue *varg = &vargs[0]; + + Assert(nargs == 1); + + if (varg->type == PGBT_INT) + { + int64 i = varg->u.ival; + + setIntValue(retval, i < 0 ? -i : i); + } + else + { + double d = varg->u.dval; + + Assert(varg->type == PGBT_DOUBLE); + setDoubleValue(retval, d < 0.0 ? -d : d); + } + + return true; + } + + case PGBENCH_DEBUG: + { + PgBenchValue *varg = &vargs[0]; + + Assert(nargs == 1); + + fprintf(stderr, "debug(script=%d,command=%d): ", + st->use_file, st->command + 1); + + if (varg->type == PGBT_NULL) + fprintf(stderr, "null\n"); + else if (varg->type == PGBT_BOOLEAN) + fprintf(stderr, "boolean %s\n", varg->u.bval ? "true" : "false"); + else if (varg->type == PGBT_INT) + fprintf(stderr, "int " INT64_FORMAT "\n", varg->u.ival); + else if (varg->type == PGBT_DOUBLE) + fprintf(stderr, "double %.*g\n", DBL_DIG, varg->u.dval); + else /* internal error, unexpected type */ + Assert(0); + + *retval = *varg; + + return true; + } + + /* 1 double argument */ + case PGBENCH_DOUBLE: + case PGBENCH_SQRT: + case PGBENCH_LN: + case PGBENCH_EXP: + { + double dval; + + Assert(nargs == 1); + + if (!coerceToDouble(&vargs[0], &dval)) + return false; + + if (func == PGBENCH_SQRT) + dval = sqrt(dval); + else if (func == PGBENCH_LN) + dval = log(dval); + else if (func == PGBENCH_EXP) + dval = exp(dval); + /* else is cast: do nothing */ + + setDoubleValue(retval, dval); + return true; + } + + /* 1 int argument */ + case PGBENCH_INT: + { + int64 ival; + + Assert(nargs == 1); + + if (!coerceToInt(&vargs[0], &ival)) + return false; + + setIntValue(retval, ival); + return true; + } + + /* variable number of arguments */ + case PGBENCH_LEAST: + case PGBENCH_GREATEST: + { + bool havedouble; + int i; + + Assert(nargs >= 1); + + /* need double result if any input is double */ + havedouble = false; + for (i = 0; i < nargs; i++) + { + if (vargs[i].type == PGBT_DOUBLE) + { + havedouble = true; + break; + } + } + if (havedouble) + { + double extremum; + + if (!coerceToDouble(&vargs[0], &extremum)) + return false; + for (i = 1; i < nargs; i++) + { + double dval; + + if (!coerceToDouble(&vargs[i], &dval)) + return false; + if (func == PGBENCH_LEAST) + extremum = Min(extremum, dval); + else + extremum = Max(extremum, dval); + } + setDoubleValue(retval, extremum); + } + else + { + int64 extremum; + + if (!coerceToInt(&vargs[0], &extremum)) + return false; + for (i = 1; i < nargs; i++) + { + int64 ival; + + if (!coerceToInt(&vargs[i], &ival)) + return false; + if (func == PGBENCH_LEAST) + extremum = Min(extremum, ival); + else + extremum = Max(extremum, ival); + } + setIntValue(retval, extremum); + } + return true; + } + + /* random functions */ + case PGBENCH_RANDOM: + case PGBENCH_RANDOM_EXPONENTIAL: + case PGBENCH_RANDOM_GAUSSIAN: + case PGBENCH_RANDOM_ZIPFIAN: + { + int64 imin, + imax, + delta; + + Assert(nargs >= 2); + + if (!coerceToInt(&vargs[0], &imin) || + !coerceToInt(&vargs[1], &imax)) + return false; + + /* check random range */ + if (unlikely(imin > imax)) + { + pg_log_error("empty range given to random"); + return false; + } + else if (unlikely(pg_sub_s64_overflow(imax, imin, &delta) || + pg_add_s64_overflow(delta, 1, &delta))) + { + /* prevent int overflows in random functions */ + pg_log_error("random range is too large"); + return false; + } + + if (func == PGBENCH_RANDOM) + { + Assert(nargs == 2); + setIntValue(retval, getrand(&st->cs_func_rs, imin, imax)); + } + else /* gaussian & exponential */ + { + double param; + + Assert(nargs == 3); + + if (!coerceToDouble(&vargs[2], ¶m)) + return false; + + if (func == PGBENCH_RANDOM_GAUSSIAN) + { + if (param < MIN_GAUSSIAN_PARAM) + { + pg_log_error("gaussian parameter must be at least %f (not %f)", + MIN_GAUSSIAN_PARAM, param); + return false; + } + + setIntValue(retval, + getGaussianRand(&st->cs_func_rs, + imin, imax, param)); + } + else if (func == PGBENCH_RANDOM_ZIPFIAN) + { + if (param < MIN_ZIPFIAN_PARAM || param > MAX_ZIPFIAN_PARAM) + { + pg_log_error("zipfian parameter must be in range [%.3f, %.0f] (not %f)", + MIN_ZIPFIAN_PARAM, MAX_ZIPFIAN_PARAM, param); + return false; + } + + setIntValue(retval, + getZipfianRand(&st->cs_func_rs, imin, imax, param)); + } + else /* exponential */ + { + if (param <= 0.0) + { + pg_log_error("exponential parameter must be greater than zero (not %f)", + param); + return false; + } + + setIntValue(retval, + getExponentialRand(&st->cs_func_rs, + imin, imax, param)); + } + } + + return true; + } + + case PGBENCH_POW: + { + PgBenchValue *lval = &vargs[0]; + PgBenchValue *rval = &vargs[1]; + double ld, + rd; + + Assert(nargs == 2); + + if (!coerceToDouble(lval, &ld) || + !coerceToDouble(rval, &rd)) + return false; + + setDoubleValue(retval, pow(ld, rd)); + + return true; + } + + case PGBENCH_IS: + { + Assert(nargs == 2); + + /* + * note: this simple implementation is more permissive than + * SQL + */ + setBoolValue(retval, + vargs[0].type == vargs[1].type && + vargs[0].u.bval == vargs[1].u.bval); + return true; + } + + /* hashing */ + case PGBENCH_HASH_FNV1A: + case PGBENCH_HASH_MURMUR2: + { + int64 val, + seed; + + Assert(nargs == 2); + + if (!coerceToInt(&vargs[0], &val) || + !coerceToInt(&vargs[1], &seed)) + return false; + + if (func == PGBENCH_HASH_MURMUR2) + setIntValue(retval, getHashMurmur2(val, seed)); + else if (func == PGBENCH_HASH_FNV1A) + setIntValue(retval, getHashFnv1a(val, seed)); + else + /* cannot get here */ + Assert(0); + + return true; + } + + case PGBENCH_PERMUTE: + { + int64 val, + size, + seed; + + Assert(nargs == 3); + + if (!coerceToInt(&vargs[0], &val) || + !coerceToInt(&vargs[1], &size) || + !coerceToInt(&vargs[2], &seed)) + return false; + + if (size <= 0) + { + pg_log_error("permute size parameter must be greater than zero"); + return false; + } + + setIntValue(retval, permute(val, size, seed)); + return true; + } + + default: + /* cannot get here */ + Assert(0); + /* dead code to avoid a compiler warning */ + return false; + } +} + +/* evaluate some function */ +static bool +evalFunc(CState *st, + PgBenchFunction func, PgBenchExprLink *args, PgBenchValue *retval) +{ + if (isLazyFunc(func)) + return evalLazyFunc(st, func, args, retval); + else + return evalStandardFunc(st, func, args, retval); +} + +/* + * Recursive evaluation of an expression in a pgbench script + * using the current state of variables. + * Returns whether the evaluation was ok, + * the value itself is returned through the retval pointer. + */ +bool +evaluateExpr(CState *st, PgBenchExpr *expr, PgBenchValue *retval) +{ + switch (expr->etype) + { + case ENODE_CONSTANT: + { + *retval = expr->u.constant; + return true; + } + + case ENODE_VARIABLE: + { + Variable *var; + + if ((var = lookupVariable(&st->variables, expr->u.variable.varname)) == NULL) + { + pg_log_error("undefined variable \"%s\"", expr->u.variable.varname); + return false; + } + + if (!makeVariableValue(var)) + return false; + + *retval = var->value; + return true; + } + + case ENODE_FUNCTION: + return evalFunc(st, + expr->u.function.function, + expr->u.function.args, + retval); + + default: + /* internal error which should never occur */ + pg_fatal("unexpected enode type in evaluation: %d", expr->etype); + } +} + +/* + * Convert command name to meta-command enum identifier + */ +static MetaCommand +getMetaCommand(const char *cmd) +{ + MetaCommand mc; + + if (cmd == NULL) + mc = META_NONE; + else if (pg_strcasecmp(cmd, "set") == 0) + mc = META_SET; + else if (pg_strcasecmp(cmd, "setshell") == 0) + mc = META_SETSHELL; + else if (pg_strcasecmp(cmd, "shell") == 0) + mc = META_SHELL; + else if (pg_strcasecmp(cmd, "sleep") == 0) + mc = META_SLEEP; + else if (pg_strcasecmp(cmd, "if") == 0) + mc = META_IF; + else if (pg_strcasecmp(cmd, "elif") == 0) + mc = META_ELIF; + else if (pg_strcasecmp(cmd, "else") == 0) + mc = META_ELSE; + else if (pg_strcasecmp(cmd, "endif") == 0) + mc = META_ENDIF; + else if (pg_strcasecmp(cmd, "gset") == 0) + mc = META_GSET; + else if (pg_strcasecmp(cmd, "aset") == 0) + mc = META_ASET; + else if (pg_strcasecmp(cmd, "startpipeline") == 0) + mc = META_STARTPIPELINE; + else if (pg_strcasecmp(cmd, "syncpipeline") == 0) + mc = META_SYNCPIPELINE; + else if (pg_strcasecmp(cmd, "endpipeline") == 0) + mc = META_ENDPIPELINE; + else + mc = META_NONE; + return mc; +} + +/* return a script number with a weighted choice. */ +int +chooseScript(TState *thread) +{ + int i = 0; + int64 w; + + if (num_scripts == 1) + return 0; + + w = getrand(&thread->ts_choose_rs, 0, total_weight - 1); + do + { + w -= sql_script[i++].weight; + } while (w >= 0); + + return i - 1; +} + +/* + * Replace :param with $n throughout the command's SQL text, which + * is a modifiable string in cmd->lines. + */ +static bool +parseQuery(Command *cmd) +{ + char *sql, + *p; + + cmd->argc = 1; + + p = sql = pg_strdup(cmd->lines.data); + while ((p = strchr(p, ':')) != NULL) + { + char var[13]; + char *name; + int eaten; + + name = parseVariable(p, &eaten); + if (name == NULL) + { + while (*p == ':') + { + p++; + } + continue; + } + + /* + * cmd->argv[0] is the SQL statement itself, so the max number of + * arguments is one less than MAX_ARGS + */ + if (cmd->argc >= MAX_ARGS) + { + pg_log_error("statement has too many arguments (maximum is %d): %s", + MAX_ARGS - 1, cmd->lines.data); + pg_free(name); + return false; + } + + sprintf(var, "$%d", cmd->argc); + p = replaceVariable(&sql, p, eaten, var); + + cmd->argv[cmd->argc] = name; + cmd->argc++; + } + + Assert(cmd->argv[0] == NULL); + cmd->argv[0] = sql; + return true; +} + +/* + * syntax error while parsing a script (in practice, while parsing a + * backslash command, because we don't detect syntax errors in SQL) + * + * source: source of script (filename or builtin-script ID) + * lineno: line number within script (count from 1) + * line: whole line of backslash command, if available + * command: backslash command name, if available + * msg: the actual error message + * more: optional extra message + * column: zero-based column number, or -1 if unknown + */ +void +syntax_error(const char *source, int lineno, + const char *line, const char *command, + const char *msg, const char *more, int column) +{ + PQExpBufferData buf; + + initPQExpBuffer(&buf); + + printfPQExpBuffer(&buf, "%s:%d: %s", source, lineno, msg); + if (more != NULL) + appendPQExpBuffer(&buf, " (%s)", more); + if (column >= 0 && line == NULL) + appendPQExpBuffer(&buf, " at column %d", column + 1); + if (command != NULL) + appendPQExpBuffer(&buf, " in command \"%s\"", command); + + pg_log_error("%s", buf.data); + + termPQExpBuffer(&buf); + + if (line != NULL) + { + fprintf(stderr, "%s\n", line); + if (column >= 0) + fprintf(stderr, "%*c error found here\n", column + 1, '^'); + } + + exit(1); +} + +/* + * Return a pointer to the start of the SQL command, after skipping over + * whitespace and "--" comments. + * If the end of the string is reached, return NULL. + */ +static char * +skip_sql_comments(char *sql_command) +{ + char *p = sql_command; + + /* Skip any leading whitespace, as well as "--" style comments */ + for (;;) + { + if (isspace((unsigned char) *p)) + p++; + else if (strncmp(p, "--", 2) == 0) + { + p = strchr(p, '\n'); + if (p == NULL) + return NULL; + p++; + } + else + break; + } + + /* NULL if there's nothing but whitespace and comments */ + if (*p == '\0') + return NULL; + + return p; +} + +/* + * Parse a SQL command; return a Command struct, or NULL if it's a comment + * + * On entry, psqlscan.l has collected the command into "buf", so we don't + * really need to do much here except check for comments and set up a Command + * struct. + */ +static Command * +create_sql_command(PQExpBuffer buf) +{ + Command *my_command; + char *p = skip_sql_comments(buf->data); + + if (p == NULL) + return NULL; + + /* Allocate and initialize Command structure */ + my_command = pg_malloc0_object(Command); + initPQExpBuffer(&my_command->lines); + appendPQExpBufferStr(&my_command->lines, p); + my_command->first_line = NULL; /* this is set later */ + my_command->type = SQL_COMMAND; + my_command->meta = META_NONE; + my_command->argc = 0; + my_command->retries = 0; + my_command->failures = 0; + memset(my_command->argv, 0, sizeof(my_command->argv)); + my_command->varprefix = NULL; /* allocated later, if needed */ + my_command->expr = NULL; + initSimpleStats(&my_command->stats); + my_command->prepname = NULL; /* set later, if needed */ + + return my_command; +} + +/* Free a Command structure and associated data */ +static void +free_command(Command *command) +{ + termPQExpBuffer(&command->lines); + pg_free(command->first_line); + for (int i = 0; i < command->argc; i++) + pg_free(command->argv[i]); + pg_free(command->varprefix); + + /* + * It should also free expr recursively, but this is currently not needed + * as only gset commands (which do not have an expression) are freed. + */ + pg_free(command); +} + +/* + * Once an SQL command is fully parsed, possibly by accumulating several + * parts, complete other fields of the Command structure. + */ +void +postprocess_sql_command(Command *my_command) +{ + char buffer[128]; + static int prepnum = 0; + + Assert(my_command->type == SQL_COMMAND); + + /* Save the first line for error display. */ + strlcpy(buffer, my_command->lines.data, sizeof(buffer)); + buffer[strcspn(buffer, "\n\r")] = '\0'; + my_command->first_line = pg_strdup(buffer); + + /* Parse query and generate prepared statement name, if necessary */ + switch (querymode) + { + case QUERY_SIMPLE: + my_command->argv[0] = my_command->lines.data; + my_command->argc++; + break; + case QUERY_PREPARED: + my_command->prepname = psprintf("P_%d", prepnum++); + pg_fallthrough; + case QUERY_EXTENDED: + if (!parseQuery(my_command)) + exit(1); + break; + default: + exit(1); + } +} + +/* + * Parse a backslash command; return a Command struct, or NULL if comment + * + * At call, we have scanned only the initial backslash. + */ +static Command * +process_backslash_command(PsqlScanState sstate, const char *source, + int lineno, int start_offset) +{ + Command *my_command; + PQExpBufferData word_buf; + int word_offset; + int offsets[MAX_ARGS]; /* offsets of argument words */ + int j; + + initPQExpBuffer(&word_buf); + + /* Collect first word of command */ + if (!expr_lex_one_word(sstate, &word_buf, &word_offset)) + { + termPQExpBuffer(&word_buf); + return NULL; + } + + /* Allocate and initialize Command structure */ + my_command = pg_malloc0_object(Command); + my_command->type = META_COMMAND; + my_command->argc = 0; + initSimpleStats(&my_command->stats); + + /* Save first word (command name) */ + j = 0; + offsets[j] = word_offset; + my_command->argv[j++] = pg_strdup(word_buf.data); + my_command->argc++; + + /* ... and convert it to enum form */ + my_command->meta = getMetaCommand(my_command->argv[0]); + + if (my_command->meta == META_SET || + my_command->meta == META_IF || + my_command->meta == META_ELIF) + { + yyscan_t yyscanner; + + /* For \set, collect var name */ + if (my_command->meta == META_SET) + { + if (!expr_lex_one_word(sstate, &word_buf, &word_offset)) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "missing argument", NULL, -1); + + offsets[j] = word_offset; + my_command->argv[j++] = pg_strdup(word_buf.data); + my_command->argc++; + } + + /* then for all parse the expression */ + yyscanner = expr_scanner_init(sstate, source, lineno, start_offset, + my_command->argv[0]); + + if (expr_yyparse(&my_command->expr, yyscanner) != 0) + { + /* dead code: exit done from syntax_error called by yyerror */ + exit(1); + } + + /* Save line, trimming any trailing newline */ + my_command->first_line = + expr_scanner_get_substring(sstate, + start_offset, + true); + + expr_scanner_finish(yyscanner); + + termPQExpBuffer(&word_buf); + + return my_command; + } + + /* For all other commands, collect remaining words. */ + while (expr_lex_one_word(sstate, &word_buf, &word_offset)) + { + /* + * my_command->argv[0] is the command itself, so the max number of + * arguments is one less than MAX_ARGS + */ + if (j >= MAX_ARGS) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "too many arguments", NULL, -1); + + offsets[j] = word_offset; + my_command->argv[j++] = pg_strdup(word_buf.data); + my_command->argc++; + } + + /* Save line, trimming any trailing newline */ + my_command->first_line = + expr_scanner_get_substring(sstate, + start_offset, + true); + + if (my_command->meta == META_SLEEP) + { + if (my_command->argc < 2) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "missing argument", NULL, -1); + + if (my_command->argc > 3) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "too many arguments", NULL, + offsets[3] - start_offset); + + /* + * Split argument into number and unit to allow "sleep 1ms" etc. We + * don't have to terminate the number argument with null because it + * will be parsed with atoi, which ignores trailing non-digit + * characters. + */ + 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; + + /* Eat all digits */ + while (*c && isdigit((unsigned char) *c)) + c++; + + if (*c) + { + 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); + } + } + } + + if (my_command->argc == 3) + { + if (pg_strcasecmp(my_command->argv[2], "us") != 0 && + pg_strcasecmp(my_command->argv[2], "ms") != 0 && + pg_strcasecmp(my_command->argv[2], "s") != 0) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "unrecognized time unit, must be us, ms or s", + my_command->argv[2], offsets[2] - start_offset); + } + } + else if (my_command->meta == META_SETSHELL) + { + if (my_command->argc < 3) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "missing argument", NULL, -1); + } + else if (my_command->meta == META_SHELL) + { + if (my_command->argc < 2) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "missing command", NULL, -1); + } + else if (my_command->meta == META_ELSE || my_command->meta == META_ENDIF || + my_command->meta == META_STARTPIPELINE || + my_command->meta == META_ENDPIPELINE || + my_command->meta == META_SYNCPIPELINE) + { + if (my_command->argc != 1) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "unexpected argument", NULL, -1); + } + else if (my_command->meta == META_GSET || my_command->meta == META_ASET) + { + if (my_command->argc > 2) + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "too many arguments", NULL, -1); + } + else + { + /* my_command->meta == META_NONE */ + syntax_error(source, lineno, my_command->first_line, my_command->argv[0], + "invalid command", NULL, -1); + } + + termPQExpBuffer(&word_buf); + + return my_command; +} + +static void +ConditionError(const char *desc, int cmdn, const char *msg) +{ + pg_fatal("condition error in script \"%s\" command %d: %s", + desc, cmdn, msg); +} + +/* + * Partial evaluation of conditionals before recording and running the script. + */ +static void +CheckConditional(const ParsedScript *ps) +{ + /* statically check conditional structure */ + ConditionalStack cs = conditional_stack_create(); + int i; + + for (i = 0; ps->commands[i] != NULL; i++) + { + Command *cmd = ps->commands[i]; + + if (cmd->type == META_COMMAND) + { + switch (cmd->meta) + { + case META_IF: + conditional_stack_push(cs, IFSTATE_FALSE); + break; + case META_ELIF: + if (conditional_stack_empty(cs)) + ConditionError(ps->desc, i + 1, "\\elif without matching \\if"); + if (conditional_stack_peek(cs) == IFSTATE_ELSE_FALSE) + ConditionError(ps->desc, i + 1, "\\elif after \\else"); + break; + case META_ELSE: + if (conditional_stack_empty(cs)) + ConditionError(ps->desc, i + 1, "\\else without matching \\if"); + if (conditional_stack_peek(cs) == IFSTATE_ELSE_FALSE) + ConditionError(ps->desc, i + 1, "\\else after \\else"); + conditional_stack_poke(cs, IFSTATE_ELSE_FALSE); + break; + case META_ENDIF: + if (!conditional_stack_pop(cs)) + ConditionError(ps->desc, i + 1, "\\endif without matching \\if"); + break; + default: + /* ignore anything else... */ + break; + } + } + } + if (!conditional_stack_empty(cs)) + ConditionError(ps->desc, i + 1, "\\if without matching \\endif"); + conditional_stack_destroy(cs); +} + +/* + * Parse a script (either the contents of a file, or a built-in script) + * and add it to the list of scripts. + */ +static void +ParseScript(const char *script, const char *desc, int weight) +{ + ParsedScript ps; + PsqlScanState sstate; + PQExpBufferData line_buf; + int alloc_num; + int index; + +#define COMMANDS_ALLOC_NUM 128 + alloc_num = COMMANDS_ALLOC_NUM; + + /* Initialize all fields of ps */ + ps.desc = desc; + ps.weight = weight; + ps.commands = pg_malloc_array(Command *, alloc_num); + initStats(&ps.stats, 0); + + /* Prepare to parse script */ + sstate = psql_scan_create(&pgbench_callbacks); + + /* + * Ideally, we'd scan scripts using the encoding and stdstrings settings + * we get from a DB connection. However, without major rearrangement of + * pgbench's argument parsing, we can't have a DB connection at the time + * we parse scripts. Using SQL_ASCII (encoding 0) should work well enough + * with any backend-safe encoding, though conceivably we could be fooled + * if a script file uses a client-only encoding. We also assume that + * stdstrings should be true, which is a bit riskier. + */ + psql_scan_setup(sstate, script, strlen(script), 0, true); + + initPQExpBuffer(&line_buf); + + index = 0; + + for (;;) + { + PsqlScanResult sr; + promptStatus_t prompt; + Command *command = NULL; + + resetPQExpBuffer(&line_buf); + + sr = psql_scan(sstate, &line_buf, &prompt); + + /* If we collected a new SQL command, process that */ + command = create_sql_command(&line_buf); + + /* store new command */ + if (command) + ps.commands[index++] = command; + + /* If we reached a backslash, process that */ + if (sr == PSCAN_BACKSLASH) + { + int lineno; + int start_offset; + + /* Capture location of the backslash */ + psql_scan_get_location(sstate, &lineno, &start_offset); + start_offset--; + + command = process_backslash_command(sstate, desc, + lineno, start_offset); + + if (command) + { + /* + * If this is gset or aset, merge into the preceding command. + * (We don't use a command slot in this case). + */ + if (command->meta == META_GSET || command->meta == META_ASET) + { + Command *cmd; + + if (index == 0) + syntax_error(desc, lineno, NULL, NULL, + "\\gset must follow an SQL command", + NULL, -1); + + cmd = ps.commands[index - 1]; + + if (cmd->type != SQL_COMMAND || + cmd->varprefix != NULL) + syntax_error(desc, lineno, NULL, NULL, + "\\gset must follow an SQL command", + cmd->first_line, -1); + + /* get variable prefix */ + if (command->argc <= 1 || command->argv[1][0] == '\0') + cmd->varprefix = pg_strdup(""); + else + cmd->varprefix = pg_strdup(command->argv[1]); + + /* update the sql command meta */ + cmd->meta = command->meta; + + /* cleanup unused command */ + free_command(command); + + continue; + } + + /* Attach any other backslash command as a new command */ + ps.commands[index++] = command; + } + } + + /* + * Since we used a command slot, allocate more if needed. Note we + * always allocate one more in order to accommodate the NULL + * terminator below. + */ + if (index >= alloc_num) + { + alloc_num += COMMANDS_ALLOC_NUM; + ps.commands = (Command **) + pg_realloc_array(ps.commands, Command *, alloc_num); + } + + /* Done if we reached EOF */ + if (sr == PSCAN_INCOMPLETE || sr == PSCAN_EOL) + break; + } + + ps.commands[index] = NULL; + + addScript(&ps); + + termPQExpBuffer(&line_buf); + psql_scan_finish(sstate); + psql_scan_destroy(sstate); +} + +/* + * Read the entire contents of file fd, and return it in a malloc'd buffer. + * + * The buffer will typically be larger than necessary, but we don't care + * in this program, because we'll free it as soon as we've parsed the script. + */ +static char * +read_file_contents(FILE *fd) +{ + char *buf; + size_t buflen = BUFSIZ; + size_t used = 0; + + buf = (char *) pg_malloc(buflen); + + for (;;) + { + size_t nread; + + nread = fread(buf + used, 1, BUFSIZ, fd); + used += nread; + /* If fread() read less than requested, must be EOF or error */ + if (nread < BUFSIZ) + break; + /* Enlarge buf so we can read some more */ + buflen += BUFSIZ; + buf = (char *) pg_realloc(buf, buflen); + } + /* There is surely room for a terminator */ + buf[used] = '\0'; + + return buf; +} + +/* + * Given a file name, read it and add its script to the list. + * "-" means to read stdin. + * NB: filename must be storage that won't disappear. + */ +void +process_file(const char *filename, int weight) +{ + FILE *fd; + char *buf; + + /* Slurp the file contents into "buf" */ + if (strcmp(filename, "-") == 0) + fd = stdin; + else if ((fd = fopen(filename, "r")) == NULL) + pg_fatal("could not open file \"%s\": %m", filename); + + buf = read_file_contents(fd); + + if (ferror(fd)) + pg_fatal("could not read file \"%s\": %m", filename); + + if (fd != stdin) + fclose(fd); + + ParseScript(buf, filename, weight); + + free(buf); +} + +/* Parse the given builtin script and add it to the list. */ +void +process_builtin(const BuiltinScript *bi, int weight) +{ + ParseScript(bi->script, bi->desc, weight); +} + +/* show available builtin scripts */ +void +listAvailableScripts(void) +{ + fprintf(stderr, "Available builtin scripts:\n"); + for (size_t i = 0; i < lengthof(builtin_script); i++) + fprintf(stderr, " %13s: %s\n", builtin_script[i].name, builtin_script[i].desc); + fprintf(stderr, "\n"); +} + +/* return builtin script "name" if unambiguous, fails if not found */ +const BuiltinScript * +findBuiltin(const char *name) +{ + int found = 0, + len = strlen(name); + const BuiltinScript *result = NULL; + + for (size_t i = 0; i < lengthof(builtin_script); i++) + { + if (strncmp(builtin_script[i].name, name, len) == 0) + { + result = &builtin_script[i]; + found++; + } + } + + /* ok, unambiguous result */ + if (found == 1) + return result; + + /* error cases */ + if (found == 0) + pg_log_error("no builtin script found for name \"%s\"", name); + else /* found > 1 */ + pg_log_error("ambiguous builtin name: %d builtin scripts found for prefix \"%s\"", found, name); + + listAvailableScripts(); + exit(1); +} + +/* + * Determine the weight specification from a script option (-b, -f), if any, + * and return it as an integer (1 is returned if there's no weight). The + * script name is returned in *script as a malloc'd string. + */ +int +parseScriptWeight(const char *option, char **script) +{ + const char *sep; + int weight; + + if ((sep = strrchr(option, WSEP))) + { + int namelen = sep - option; + long wtmp; + char *badp; + + /* generate the script name */ + *script = pg_malloc(namelen + 1); + strncpy(*script, option, namelen); + (*script)[namelen] = '\0'; + + /* process digits of the weight spec */ + errno = 0; + wtmp = strtol(sep + 1, &badp, 10); + if (errno != 0 || badp == sep + 1 || *badp != '\0') + pg_fatal("invalid weight specification: %s", sep); + if (wtmp > INT_MAX || wtmp < 0) + pg_fatal("weight specification out of range (0 .. %d): %ld", + INT_MAX, wtmp); + weight = wtmp; + } + else + { + *script = pg_strdup(option); + weight = 1; + } + + return weight; +} + +/* append a script to the list of scripts to process */ +static void +addScript(const ParsedScript *script) +{ + if (script->commands == NULL || script->commands[0] == NULL) + pg_fatal("empty command list for script \"%s\"", script->desc); + + if (num_scripts >= MAX_SCRIPTS) + pg_fatal("at most %d SQL scripts are allowed", MAX_SCRIPTS); + + CheckConditional(script); + + sql_script[num_scripts] = *script; + num_scripts++; +} diff --git a/src/bin/pgbench/script.h b/src/bin/pgbench/script.h new file mode 100644 index 00000000000..69b8f0aa9b2 --- /dev/null +++ b/src/bin/pgbench/script.h @@ -0,0 +1,170 @@ +/*------------------------------------------------------------------------- + * + * script.h + * Script parsing, built-in scripts, and expression evaluation for pgbench + * + * Portions Copyright (c) 1996-2026, PostgreSQL Global Development Group + * Portions Copyright (c) 1994, Regents of the University of California + * + * src/bin/pgbench/script.h + * + *------------------------------------------------------------------------- + */ +#ifndef SCRIPT_H +#define SCRIPT_H + +#include "common/pg_prng.h" +#include "pgbench.h" +#include "pqexpbuffer.h" +#include "stats.h" +#include "variable.h" + +/* X/Open (XSI) requires to provide M_PI, but core POSIX does not */ +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +/* + * Hashing constants + */ +#define FNV_PRIME UINT64CONST(0x100000001b3) +#define FNV_OFFSET_BASIS UINT64CONST(0xcbf29ce484222325) +#define MM2_MUL UINT64CONST(0xc6a4a7935bd1e995) +#define MM2_MUL_TIMES_8 UINT64CONST(0x35253c9ade8f4ca8) +#define MM2_ROT 47 + +#define MIN_GAUSSIAN_PARAM 2.0 /* minimum parameter for gauss */ + +#define MIN_ZIPFIAN_PARAM 1.001 /* minimum parameter for zipfian */ +#define MAX_ZIPFIAN_PARAM 1000.0 /* maximum parameter for zipfian */ + +#define nbranches 1 /* Makes little sense to change this. Change + * -s instead */ +#define ntellers 10 +#define naccounts 100000 + +/* + * The scale factor at/beyond which 32bit integers are incapable of storing + * 64bit values. + * + * Although the actual threshold is 21474, we use 20000 because it is easier to + * document and remember, and isn't that far away from the real threshold. + */ +#define SCALE_32BIT_THRESHOLD 20000 + +#define WSEP '@' /* weight separator */ + +#define MAX_SCRIPTS 128 /* max number of SQL scripts allowed */ + +/* + * queries read from files + */ +#define SQL_COMMAND 1 +#define META_COMMAND 2 + +/* + * max number of backslash command arguments or SQL variables, + * including the command or SQL statement itself + */ +#define MAX_ARGS 256 + +typedef enum MetaCommand +{ + META_NONE, /* not a known meta-command */ + META_SET, /* \set */ + META_SETSHELL, /* \setshell */ + META_SHELL, /* \shell */ + META_SLEEP, /* \sleep */ + META_GSET, /* \gset */ + META_ASET, /* \aset */ + META_IF, /* \if */ + META_ELIF, /* \elif */ + META_ELSE, /* \else */ + META_ENDIF, /* \endif */ + META_STARTPIPELINE, /* \startpipeline */ + META_SYNCPIPELINE, /* \syncpipeline */ + META_ENDPIPELINE, /* \endpipeline */ +} MetaCommand; + +typedef enum QueryMode +{ + QUERY_SIMPLE, /* simple query */ + QUERY_EXTENDED, /* extended query */ + QUERY_PREPARED, /* extended query with prepared statements */ + NUM_QUERYMODE +} QueryMode; + +extern QueryMode querymode; +extern const char *const QUERYMODE[]; + +/* + * struct Command represents one command in a script. + * + * lines The raw, possibly multi-line command text. Variable substitution + * not applied. + * first_line A short, single-line extract of 'lines', for error reporting. + * type SQL_COMMAND or META_COMMAND + * meta The type of meta-command, with META_NONE/GSET/ASET if command + * is SQL. + * argc Number of arguments of the command, 0 if not yet processed. + * argv Command arguments, the first of which is the command or SQL + * string itself. For SQL commands, after post-processing + * argv[0] is the same as 'lines' with variables substituted. + * prepname The name that this command is prepared under, in prepare mode + * varprefix SQL commands terminated with \gset or \aset have this set + * to a non NULL value. If nonempty, it's used to prefix the + * variable name that receives the value. + * aset do gset on all possible queries of a combined query (\;). + * expr Parsed expression, if needed. + * stats Time spent in this command. + * retries Number of retries after a serialization or deadlock error in the + * current command. + * failures Number of errors in the current command that were not retried. + */ +typedef struct Command +{ + PQExpBufferData lines; + char *first_line; + int type; + MetaCommand meta; + int argc; + char *argv[MAX_ARGS]; + char *prepname; + char *varprefix; + PgBenchExpr *expr; + SimpleStats stats; + int64 retries; + int64 failures; +} Command; + +typedef struct ParsedScript +{ + const char *desc; /* script descriptor (eg, file name) */ + int weight; /* selection weight */ + Command **commands; /* NULL-terminated array of Commands */ + StatsData stats; /* total time spent in script */ +} ParsedScript; + +extern ParsedScript sql_script[MAX_SCRIPTS]; +extern int num_scripts; +extern int64 total_weight; + +/* Builtin test scripts */ +typedef struct BuiltinScript +{ + const char *name; /* very short name for -b ... */ + const char *desc; /* short description */ + const char *script; /* actual pgbench script */ +} BuiltinScript; + +extern bool evaluateExpr(CState *st, PgBenchExpr *expr, PgBenchValue *retval); +extern int chooseScript(TState *thread); + +extern void postprocess_sql_command(Command *my_command); +extern void process_file(const char *filename, int weight); +extern void process_builtin(const BuiltinScript *bi, int weight); +extern void listAvailableScripts(void); +extern const BuiltinScript *findBuiltin(const char *tb); +extern int parseScriptWeight(const char *option, char **script); + +#endif /* SCRIPT_H */ -- 2.56.0.385.gd3acb90ef8-goog