From 2164409cb62f3dc3904822f56a8cc7cd3a4e4e18 Mon Sep 17 00:00:00 2001 From: Hannu Krosing Date: Sun, 4 Oct 2026 11:24:49 +0000 Subject: [PATCH v4 12/12] Refactor pgbench commands: introduce pluggable meta-command dispatch Refactor meta-command execution and result processing in commands.c/h and engine.c: - Introduce CommandHandler callback signature and CommandDescriptor table (command_handlers[]) in commands.c/h, decomposing the monolithic executeMetaCommand() switch into per-command handler functions (handle_cmd_sleep, handle_cmd_set, handle_cmd_if, handle_cmd_elif, handle_cmd_else, handle_cmd_endif, handle_cmd_setshell, handle_cmd_shell, handle_cmd_startpipeline, handle_cmd_syncpipeline, handle_cmd_endpipeline). - Extract skipConditionalCommands(CState *st) into commands.c/h to handle skipping inactive conditional branches, replacing the inline loop in advanceConnectionState() (engine.c). - Extract processGSetResult() into commands.c/h to handle storing query results for \gset and \aset commands, replacing the inline PGRES_TUPLES_OK block in readCommandResponse() (engine.c). --- src/bin/pgbench/commands.c | 554 ++++++++++++++++++++++++++----------- src/bin/pgbench/commands.h | 59 +++- src/bin/pgbench/engine.c | 143 +--------- 3 files changed, 441 insertions(+), 315 deletions(-) diff --git a/src/bin/pgbench/commands.c b/src/bin/pgbench/commands.c index ab12880bb9b..0316b9d1142 100644 --- a/src/bin/pgbench/commands.c +++ b/src/bin/pgbench/commands.c @@ -23,11 +23,13 @@ #include "stats.h" #include "variable.h" +#define SHELL_COMMAND_SIZE 256 /* maximum size allowed for shell 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 +bool runShellCommand(Variables *variables, char *variable, char **argv, int argc) { char command[SHELL_COMMAND_SIZE]; @@ -162,7 +164,7 @@ commandError(CState *st, const char *message) * Allocate space for CState->prepared: we need one boolean for each command * of each script. */ -static void +void allocCStatePrepared(CState *st) { Assert(st->prepared == NULL); @@ -215,7 +217,7 @@ prepareCommand(CState *st, int command_num) * 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 +void prepareCommandsInPipeline(CState *st) { int j; @@ -251,7 +253,7 @@ prepareCommandsInPipeline(CState *st) * 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 +bool evaluateSleep(Variables *variables, int argc, char **argv, int *usecs) { char *var; @@ -293,224 +295,440 @@ evaluateSleep(Variables *variables, int argc, char **argv, int *usecs) } /* - * 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. + * Meta-command Handlers */ + ConnectionStateEnum -executeMetaCommand(CState *st, pg_time_usec_t *now) +handle_cmd_sleep(CState *st, Command *cmd, pg_time_usec_t *now) { - Command *command = sql_script[st->use_file].commands[st->command]; - int argc; - char **argv; + int usec; - Assert(command != NULL && command->type == META_COMMAND); + /* + * 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, cmd->argc, cmd->argv, &usec)) + { + commandFailed(st, "sleep", "execution of meta-command failed"); + return CSTATE_ABORTED; + } - argc = command->argc; - argv = command->argv; + pg_time_now_lazy(now); + st->sleep_until = (*now) + usec; + return CSTATE_SLEEP; +} - if (unlikely(__pg_log_level <= PG_LOG_DEBUG)) +ConnectionStateEnum +handle_cmd_set(CState *st, Command *cmd, pg_time_usec_t *now) +{ + PgBenchExpr *expr = cmd->expr; + PgBenchValue result; + + if (!evaluateExpr(st, expr, &result)) { - PQExpBufferData buf; + commandFailed(st, cmd->argv[0], "evaluation of meta-command failed"); + return CSTATE_ABORTED; + } - initPQExpBuffer(&buf); + if (!putVariableValue(&st->variables, cmd->argv[0], cmd->argv[1], &result)) + { + commandFailed(st, "set", "assignment of meta-command failed"); + return CSTATE_ABORTED; + } - printfPQExpBuffer(&buf, "client %d executing \\%s", st->id, argv[0]); - for (int i = 1; i < argc; i++) - appendPQExpBuffer(&buf, " %s", argv[i]); + return CSTATE_END_COMMAND; +} - pg_log_debug("%s", buf.data); +ConnectionStateEnum +handle_cmd_if(CState *st, Command *cmd, pg_time_usec_t *now) +{ + PgBenchExpr *expr = cmd->expr; + PgBenchValue result; + bool cond; - termPQExpBuffer(&buf); + if (!evaluateExpr(st, expr, &result)) + { + commandFailed(st, cmd->argv[0], "evaluation of meta-command failed"); + return CSTATE_ABORTED; } - if (command->meta == META_SLEEP) - { - int usec; + cond = valueTruth(&result); + conditional_stack_push(st->cstack, cond ? IFSTATE_TRUE : IFSTATE_FALSE); + return CSTATE_END_COMMAND; +} - /* - * 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; - } +ConnectionStateEnum +handle_cmd_elif(CState *st, Command *cmd, pg_time_usec_t *now) +{ + PgBenchExpr *expr = cmd->expr; + PgBenchValue result; + bool cond; - pg_time_now_lazy(now); - st->sleep_until = (*now) + usec; - return CSTATE_SLEEP; + 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; } - else if (command->meta == META_SET) + + if (!evaluateExpr(st, expr, &result)) { - PgBenchExpr *expr = command->expr; - PgBenchValue result; + commandFailed(st, cmd->argv[0], "evaluation of meta-command failed"); + return CSTATE_ABORTED; + } - 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); + return CSTATE_END_COMMAND; +} - 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) +ConnectionStateEnum +handle_cmd_else(CState *st, Command *cmd, pg_time_usec_t *now) +{ + switch (conditional_stack_peek(st->cstack)) { - /* backslash commands with an expression to evaluate */ - PgBenchExpr *expr = command->expr; - PgBenchValue result; - bool cond; + 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); + } + return CSTATE_END_COMMAND; +} - if (!evaluateExpr(st, expr, &result)) - { - commandFailed(st, argv[0], "evaluation of meta-command failed"); - return CSTATE_ABORTED; - } +ConnectionStateEnum +handle_cmd_endif(CState *st, Command *cmd, pg_time_usec_t *now) +{ + Assert(!conditional_stack_empty(st->cstack)); + conditional_stack_pop(st->cstack); + return CSTATE_END_COMMAND; +} - cond = valueTruth(&result); - conditional_stack_push(st->cstack, cond ? IFSTATE_TRUE : IFSTATE_FALSE); +ConnectionStateEnum +handle_cmd_setshell(CState *st, Command *cmd, pg_time_usec_t *now) +{ + if (!runShellCommand(&st->variables, cmd->argv[1], cmd->argv + 2, cmd->argc - 2)) + { + commandFailed(st, "setshell", "execution of meta-command failed"); + return CSTATE_ABORTED; } - else if (command->meta == META_ELIF) + return CSTATE_END_COMMAND; +} + +ConnectionStateEnum +handle_cmd_shell(CState *st, Command *cmd, pg_time_usec_t *now) +{ + if (!runShellCommand(&st->variables, NULL, cmd->argv + 1, cmd->argc - 1)) { - /* backslash commands with an expression to evaluate */ - PgBenchExpr *expr = command->expr; - PgBenchValue result; - bool cond; + commandFailed(st, "shell", "execution of meta-command failed"); + return CSTATE_ABORTED; + } + return CSTATE_END_COMMAND; +} - 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; - } +ConnectionStateEnum +handle_cmd_startpipeline(CState *st, Command *cmd, pg_time_usec_t *now) +{ + /* + * 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 (!evaluateExpr(st, expr, &result)) - { - commandFailed(st, argv[0], "evaluation of meta-command failed"); - 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); - cond = valueTruth(&result); - Assert(conditional_stack_peek(st->cstack) == IFSTATE_FALSE); - conditional_stack_poke(st->cstack, cond ? IFSTATE_TRUE : IFSTATE_FALSE); + if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) + { + commandFailed(st, "startpipeline", "already in pipeline mode"); + return CSTATE_ABORTED; } - else if (command->meta == META_ELSE) + if (PQenterPipelineMode(st->con) == 0) { - 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); - } + commandFailed(st, "startpipeline", "failed to enter pipeline mode"); + return CSTATE_ABORTED; } - else if (command->meta == META_ENDIF) + return CSTATE_END_COMMAND; +} + +ConnectionStateEnum +handle_cmd_syncpipeline(CState *st, Command *cmd, pg_time_usec_t *now) +{ + if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) { - Assert(!conditional_stack_empty(st->cstack)); - conditional_stack_pop(st->cstack); + commandFailed(st, "syncpipeline", "not in pipeline mode"); + return CSTATE_ABORTED; } - else if (command->meta == META_SETSHELL) + if (PQsendPipelineSync(st->con) == 0) { - if (!runShellCommand(&st->variables, argv[1], argv + 2, argc - 2)) - { - commandFailed(st, "setshell", "execution of meta-command failed"); - return CSTATE_ABORTED; - } + commandFailed(st, "syncpipeline", "failed to send a pipeline sync"); + return CSTATE_ABORTED; } - else if (command->meta == META_SHELL) + st->num_syncs++; + return CSTATE_END_COMMAND; +} + +ConnectionStateEnum +handle_cmd_endpipeline(CState *st, Command *cmd, pg_time_usec_t *now) +{ + if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) { - if (!runShellCommand(&st->variables, NULL, argv + 1, argc - 1)) - { - commandFailed(st, "shell", "execution of meta-command failed"); - return CSTATE_ABORTED; - } + commandFailed(st, "endpipeline", "not in pipeline mode"); + return CSTATE_ABORTED; } - else if (command->meta == META_STARTPIPELINE) + if (!PQpipelineSync(st->con)) { - /* - * In pipeline mode, we use a workflow based on libpq pipeline - * functions. - */ - if (querymode == QUERY_SIMPLE) + 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; +} + +/* Pluggable meta-command dispatch table */ +static const CommandDescriptor command_handlers[] = { + {"set", META_SET, handle_cmd_set}, + {"setshell", META_SETSHELL, handle_cmd_setshell}, + {"shell", META_SHELL, handle_cmd_shell}, + {"sleep", META_SLEEP, handle_cmd_sleep}, + {"if", META_IF, handle_cmd_if}, + {"elif", META_ELIF, handle_cmd_elif}, + {"else", META_ELSE, handle_cmd_else}, + {"endif", META_ENDIF, handle_cmd_endif}, + {"startpipeline", META_STARTPIPELINE, handle_cmd_startpipeline}, + {"syncpipeline", META_SYNCPIPELINE, handle_cmd_syncpipeline}, + {"endpipeline", META_ENDPIPELINE, handle_cmd_endpipeline}, + {NULL, META_NONE, NULL} +}; + +/* + * Execute a meta command and return the next connection state. + */ +ConnectionStateEnum +executeMetaCommand(CState *st, pg_time_usec_t *now) +{ + Command *command = sql_script[st->use_file].commands[st->command]; + int argc; + char **argv; + int i; + const CommandDescriptor *desc; + + 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 (i = 1; i < argc; i++) + appendPQExpBuffer(&buf, " %s", argv[i]); + + pg_log_debug("%s", buf.data); + + termPQExpBuffer(&buf); + } + + for (desc = command_handlers; desc->name != NULL; desc++) + { + if (desc->meta == command->meta) { - commandFailed(st, "startpipeline", "cannot use pipeline mode with the simple query protocol"); - return CSTATE_ABORTED; + ConnectionStateEnum next_state = desc->handler(st, command, now); + + if (next_state == CSTATE_END_COMMAND) + *now = 0; + + return next_state; } + } + + pg_log_error("client %d: unknown meta-command %d", st->id, (int) command->meta); + return CSTATE_ABORTED; +} + +/* + * Skip commands until an active branch or matching endif is reached. + * Implements the skipping logic for CSTATE_SKIP_COMMAND. + */ +void +skipConditionalCommands(CState *st) +{ + Command *command; + + Assert(!conditional_active(st->cstack)); + + while (true) + { + command = sql_script[st->use_file].commands[st->command]; + + /* cannot reach end of script in that state */ + Assert(command != NULL); /* - * 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 this is conditional related, update conditional + * state */ - if (querymode == QUERY_PREPARED) - prepareCommandsInPipeline(st); - - if (PQpipelineStatus(st->con) != PQ_PIPELINE_OFF) + if (command->type == META_COMMAND && + (command->meta == META_IF || + command->meta == META_ELIF || + command->meta == META_ELSE || + command->meta == META_ENDIF)) { - commandFailed(st, "startpipeline", "already in pipeline mode"); - return CSTATE_ABORTED; + 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); + } } - if (PQenterPipelineMode(st->con) == 0) + else { - commandFailed(st, "startpipeline", "failed to enter pipeline mode"); - return CSTATE_ABORTED; + /* skip and consider next */ + st->command++; } + + if (st->state != CSTATE_SKIP_COMMAND) + /* out of quick skip command loop */ + break; + } +} + +/* + * Process \gset and \aset variable storage from query result. + * Returns true if successful, false on error. + */ +bool +processGSetResult(CState *st, Command *command, PGresult *res, bool is_last, int qrynum) +{ + MetaCommand meta = command->meta; + char *varprefix = command->varprefix; + int ntuples = PQntuples(res); + const char *context = (meta == META_ASET) ? "aset" : "gset"; + int fld; + + if (meta == META_GSET && ntuples != 1) + { + pg_log_error("client %d script %d command %d query %d: expected one row, got %d", + st->id, st->use_file, st->command, qrynum, ntuples); + st->estatus = ESTATUS_META_COMMAND_ERROR; + return false; } - else if (command->meta == META_SYNCPIPELINE) + else if (meta == META_ASET && ntuples <= 0) { - 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++; + /* skip empty result under \aset */ + return true; } - else if (command->meta == META_ENDPIPELINE) + + /* store results into variables */ + for (fld = 0; fld < PQnfields(res); fld++) { - if (PQpipelineStatus(st->con) != PQ_PIPELINE_ON) + char *varname = PQfname(res, fld); + bool ok; + + if (*varprefix != '\0') + varname = psprintf("%s%s", varprefix, varname); + + /* store last row result */ + if (PQgetisnull(res, ntuples - 1, fld)) { - commandFailed(st, "endpipeline", "not in pipeline mode"); - return CSTATE_ABORTED; + PgBenchValue nullval; + + setNullValue(&nullval); + ok = putVariableValue(&st->variables, context, + varname, &nullval); } - if (!PQpipelineSync(st->con)) + else + ok = putVariable(&st->variables, context, varname, + PQgetvalue(res, ntuples - 1, fld)); + + if (!ok) { - commandFailed(st, "endpipeline", "failed to send a pipeline sync"); - return CSTATE_ABORTED; + 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; + if (*varprefix != '\0') + pfree(varname); + return false; } - 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; + if (*varprefix != '\0') + pfree(varname); + } - return CSTATE_END_COMMAND; + return true; } diff --git a/src/bin/pgbench/commands.h b/src/bin/pgbench/commands.h index 0c9bf7cfa19..05cad49df72 100644 --- a/src/bin/pgbench/commands.h +++ b/src/bin/pgbench/commands.h @@ -1,7 +1,7 @@ /*------------------------------------------------------------------------- * * commands.h - * Meta-command execution and prepared statement helpers for pgbench + * pgbench meta-commands execution and pluggable handler registry * * Portions Copyright (c) 1996-2026, PostgreSQL Global Development Group * Portions Copyright (c) 1994, Regents of the University of California @@ -10,16 +10,61 @@ * *------------------------------------------------------------------------- */ -#ifndef COMMANDS_H -#define COMMANDS_H +#ifndef PGBENCH_COMMANDS_H +#define PGBENCH_COMMANDS_H + +#include "libpq-fe.h" +#include "fe_utils/conditional.h" #include "pgbench.h" +#include "script.h" + +/* + * Pluggable Command Handler function signature. + * + * Receives the client state, the command to execute, and a timestamp pointer. + * Returns the next ConnectionStateEnum to transition the client state machine to. + */ +typedef ConnectionStateEnum (*CommandHandler)(CState *st, Command *cmd, pg_time_usec_t *now); + +/* + * Registration entry for meta-commands. + */ +typedef struct CommandDescriptor +{ + const char *name; /* e.g. "set", "sleep", "shell", "if", ... */ + MetaCommand meta; /* MetaCommand enum value */ + CommandHandler handler; /* Function executing this command */ +} CommandDescriptor; -#define SHELL_COMMAND_SIZE 256 /* maximum size allowed for shell command */ +/* Public Command Execution APIs */ +extern ConnectionStateEnum executeMetaCommand(CState *st, pg_time_usec_t *now); +extern void skipConditionalCommands(CState *st); +extern bool processGSetResult(CState *st, Command *command, PGresult *res, + bool is_last, int qrynum); +/* Helper APIs for commands */ extern void commandFailed(CState *st, const char *cmd, const char *message); extern void commandError(CState *st, const char *message); -extern void prepareCommand(CState *st, int command); -extern ConnectionStateEnum executeMetaCommand(CState *st, pg_time_usec_t *now); +extern bool evaluateSleep(Variables *variables, int argc, char **argv, int *usecs); +extern bool runShellCommand(Variables *variables, char *variable, char **argv, int argc); + +/* Handler declarations */ +extern ConnectionStateEnum handle_cmd_sleep(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_set(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_if(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_elif(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_else(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_endif(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_setshell(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_shell(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_startpipeline(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_syncpipeline(CState *st, Command *cmd, pg_time_usec_t *now); +extern ConnectionStateEnum handle_cmd_endpipeline(CState *st, Command *cmd, pg_time_usec_t *now); + +/* Preparation helpers used across commands and SQL preparation */ +extern void allocCStatePrepared(CState *st); +extern void prepareCommand(CState *st, int command_num); +extern void prepareCommandsInPipeline(CState *st); -#endif /* COMMANDS_H */ +#endif /* PGBENCH_COMMANDS_H */ diff --git a/src/bin/pgbench/engine.c b/src/bin/pgbench/engine.c index 3bbcf09ebfc..2d5d02fee47 100644 --- a/src/bin/pgbench/engine.c +++ b/src/bin/pgbench/engine.c @@ -402,58 +402,9 @@ readCommandResponse(CState *st, MetaCommand meta, char *varprefix) case PGRES_TUPLES_OK: if ((is_last && meta == META_GSET) || meta == META_ASET) { - 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; + if (!processGSetResult(st, sql_script[st->use_file].commands[st->command], + res, is_last, qrynum)) 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++) - { - 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)) - { - PgBenchValue nullval; - - setNullValue(&nullval); - ok = putVariableValue(&st->variables, context, - varname, &nullval); - } - else - ok = putVariable(&st->variables, context, varname, - PQgetvalue(res, ntuples - 1, fld)); - - if (!ok) - { - /* 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; - } - - if (*varprefix != '\0') - pfree(varname); - } } /* otherwise the result is simply thrown away by PQclear below */ break; @@ -989,95 +940,7 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg) 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; - } + skipConditionalCommands(st); break; /* -- 2.56.0.385.gd3acb90ef8-goog