# pgbench Modular Refactoring Architecture & Implementation Plan

> [!NOTE]
> **Context**: Refactoring `src/bin/pgbench/pgbench.c` (~7,700 lines) into a modular, maintainable, and extensible subsystem architecture in PostgreSQL `master`.

---

## 1. Architectural Vision & Design Decisions

### 1.1 Key Architectural Decisions (from Alignment)

1. **Subsystem Organization**: Flat structure in `src/bin/pgbench/` adhering to core PostgreSQL `src/bin/` conventions.
2. **State Management**: Explicit context objects (`PgBenchConfig`, `EngineContext`, `TState`, `CState`) passed to subsystem APIs instead of scattered file-scope globals.
3. **Command Execution**: Pluggable `CommandHandler` dispatch table with uniform function signatures and structured `CommandResult` return statuses.
4. **Variable Model**: Scoped Stack Frame Variable Store allowing nested variable lexical scopes for loops (`\for`, `\while`), transaction blocks, and procedures.
5. **Event Loop Multiplexing**: Abstracted `PgBenchPoller` / `SocketSet` interface encapsulating OS-specific multiplexers (`ppoll`, `poll`, `select`, Windows `WaitForMultipleObjects`).
6. **Concurrency Model**: Strict lockless per-thread execution (`TState` owns sockets, clients, and stats; thread-safe aggregation at reporting boundaries only).
7. **Migration Strategy**: 7-phase iterative refactoring with independently verifiable, bisectable commits.

---

## 2. Target Subsystem Breakdown

```
src/bin/pgbench/
├── pgbench.c          # Main CLI entrypoint, option parsing, runner orchestration (~400 LOC)
├── pgbench.h          # Public core types and subsystem interfaces
├── context.h          # Global configuration, engine context, thread & client state structs
│
├── stats.h / stats.c  # Latency tracking, statistics aggregation, out-of-band transaction logs
├── poller.h / poller.c# Unified cross-platform socket multiplexing (ppoll, poll, select, Win32)
├── variable.h / variable.c # Scoped Stack Frame Variable Store, lookup, creation, typing
├── script.h / script.c# Script loader, exprparse/exprscan integration, AST evaluator
├── commands.h / commands.c # Pluggable meta-command registry and handlers (\set, \sleep, \gset, etc.)
├── engine.h / engine.c# Client connection state machine, pipeline coordination, worker thread loop
└── init.h / init.c    # Schema creation, synthetic data generation, table partitioning, vacuuming
```

---

## 3. Subsystem Architecture Specifications

```
                       ┌───────────────────────────────┐
                       │           pgbench.c           │
                       │     (CLI, Main, Config)       │
                       └───────────────┬───────────────┘
                                       │
                                       ▼
                       ┌───────────────────────────────┐
                       │          context.h            │
                       │  (PgBenchConfig, TState, CState)
                       └───────────────┬───────────────┘
                                       │
         ┌─────────────────────────────┼─────────────────────────────┐
         │                             │                             │
         ▼                             ▼                             ▼
┌───────────────────┐        ┌───────────────────┐        ┌───────────────────┐
│     init.c/h      │        │     engine.c/h    │        │    stats.c/h      │
│ (Schema & Data    │        │ (Connection State │        │ (Latencies, Logs, │
│  Generation -i)   │        │  Machine & Loop)  │        │  Aggregations)    │
└───────────────────┘        └─────────┬─────────┘        └───────────────────┘
                                       │
         ┌─────────────────────────────┼─────────────────────────────┐
         │                             │                             │
         ▼                             ▼                             ▼
┌───────────────────┐        ┌───────────────────┐        ┌───────────────────┐
│    poller.c/h     │        │   commands.c/h    │        │   variable.c/h    │
│ (Socket Event     │        │ (Pluggable Meta-  │        │ (Scoped Stack     │
│  Multiplexing)    │        │  Command Registry)│        │  Variable Store)  │
└───────────────────┘        └─────────┬─────────┘        └───────────────────┘
                                       │
                                       ▼
                             ┌───────────────────┐
                             │    script.c/h     │
                             │ (AST & Expression │
                             │   Evaluation)     │
                             └───────────────────┘
```

### 3.1 Context & Configuration (`context.h`)
Encapsulates runtime state, eliminating global variables:
* `PgBenchConfig`: CLI options (scale, duration, tx count, throttle delay, sampling rates, latency limits, debug level).
* `PgBenchContext`: Global benchmark instance state, script descriptors, target weights, shared memory / sync barriers.
* `TState` (Thread Context): Per-thread state, owned client array, local PRNG sequences, local `PgBenchPoller`.
* `CState` (Client Context): Client connection, transaction state, current script AST pointer, local `VariableScopeStack`, pipeline state.

### 3.2 Scoped Variable Store (`variable.h` / `variable.c`)
Replaces flat sorted dynamic array with a scoped stack model:
```c
typedef struct VariableScope
{
    struct VariableScope   *parent;      /* Enclosing scope (NULL for global client scope) */
    int                     nvariables;
    int                     alloc;
    Variable               *vars;        /* Sorted array of variables within this scope */
} VariableScope;

typedef struct VariableScopeStack
{
    VariableScope          *current;     /* Top of scope stack */
} VariableScopeStack;

/* API */
bool var_scope_push(VariableScopeStack *stack);
bool var_scope_pop(VariableScopeStack *stack);
bool var_put_value(VariableScopeStack *stack, const char *name, const PgBenchValue *val, bool local_only);
bool var_get_value(VariableScopeStack *stack, const char *name, PgBenchValue *val);
```
* **Benefits**: Instant support for loop variables (`\for i 1 100`), local variables in stored procedures or script blocks, avoiding variable pollution.

### 3.3 Pluggable Command Handler Registry (`commands.h` / `commands.c`)
Replaces hardcoded `switch(cmd->type)` in state machines:
```c
typedef enum CommandStatus
{
    CMD_OK,         /* Command completed successfully, advance to next command */
    CMD_YIELD,      /* Yield execution back to event loop (e.g. waiting on async socket) */
    CMD_SLEEP,      /* Client scheduled for sleep */
    CMD_BRANCH,     /* Conditional or loop jump: pc modified in CState */
    CMD_ERROR       /* Runtime execution error */
} CommandStatus;

typedef struct CommandResult
{
    CommandStatus   status;
    int64           sleep_us;
    char           *error_msg;
} CommandResult;

typedef CommandResult (*CommandHandler)(CState *st, ParsedCommand *cmd);

typedef struct CommandDescriptor
{
    const char     *name;
    CommandType     type;
    CommandHandler  handler;
} CommandDescriptor;
```
* **Command Table**:
  * `\set` $\to$ `handle_cmd_set()`
  * `\sleep` $\to$ `handle_cmd_sleep()`
  * `\shell` $\to$ `handle_cmd_shell()`
  * `\gset` / `\cset` $\to$ `handle_cmd_gset()`
  * `\startpipeline` / `\endpipeline` $\to$ `handle_cmd_pipeline()`
  * `\if` / `\elif` / `\else` / `\endif` $\to$ `handle_cmd_conditional()`
  * *(Future Extension Point)*: `\for`, `\while`, `\try`, `\catch` register seamlessly without touching `engine.c`.

### 3.4 Socket Event Poller (`poller.h` / `poller.c`)
Abstracts multi-platform socket multiplexing:
```c
typedef struct PgBenchPoller PgBenchPoller;

PgBenchPoller *poller_create(int max_sockets);
void           poller_destroy(PgBenchPoller *poller);
bool           poller_add_socket(PgBenchPoller *poller, int fd, int events, void *user_data);
bool           poller_modify_socket(PgBenchPoller *poller, int fd, int events);
bool           poller_remove_socket(PgBenchPoller *poller, int fd);
int            poller_wait(PgBenchPoller *poller, int64 timeout_us, PollerEvent *events_out, int max_events);
```

### 3.5 Script AST & Expression Engine (`script.h` / `script.c`)
Encapsulates `exprparse.y`, Bison/Flex generation, script tokenization, and expression evaluation:
* `script_load_file()`, `script_load_string()`, `script_free()`
* `script_eval_expr(CState *st, PgBenchExpr *expr, PgBenchValue *out_val, char **err_msg)`

### 3.6 Connection State Machine & Engine (`engine.h` / `engine.c`)
Focuses solely on connection lifecycle and transaction flow:
* States: `CSTATE_START_TX`, `CSTATE_SEND_QUERY`, `CSTATE_WAIT_RESULT`, `CSTATE_SLEEP`, `CSTATE_FINISHED`, etc.
* `engine_step_client(CState *st)`
* `engine_thread_loop(TState *thread)`

---

## 4. Phase-by-Phase Implementation Plan

```mermaid
flowchart LR
    P1["Phase 1: stats.c/h"] --> P2["Phase 2: variable.c/h"]
    P2 --> P3["Phase 3: poller.c/h"]
    P3 --> P4["Phase 4: script.c/h"]
    P4 --> P5["Phase 5: commands.c/h"]
    P5 --> P6["Phase 6: init.c/h"]
    P6 --> P7["Phase 7: engine.c/h & pgbench.c"]
```

| Phase | Subsystem | Extracted Elements | Verification Gate |
| :--- | :--- | :--- | :--- |
| **Phase 1** | `stats.c` / `stats.h` | `SimpleStats`, `StatsData`, `accumStats()`, `mergeSimpleStats()`, out-of-band transaction logs, progress printers | `make check` (TAP: 681 tests pass) |
| **Phase 2** | `variable.c` / `variable.h` | `Variable`, `lookupCreateVariable()`, `putVariable()`, `getVariableValue()`, `VariableScopeStack` | `make check` (TAP: 681 tests pass) |
| **Phase 3** | `poller.c` / `poller.h` | `ppoll`, `poll`, `select`, Windows socket sets, event dispatch abstraction | `make check` (TAP: 681 tests pass) |
| **Phase 4** | `script.c` / `script.h` | Script loader, parser wrappers, `evalStandardFunc()`, `evalLazyFunc()`, `evaluateExpr()`, `EvalResult` | `make check` (TAP: 681 tests pass) |
| **Phase 5** | `commands.c` / `commands.h` | Pluggable `CommandHandler` table, handlers for `\set`, `\sleep`, `\shell`, `\gset`, `\startpipeline`, `\if` | `make check` (TAP: 681 tests pass) |
| **Phase 6** | `init.c` / `init.h` | `init()`, `initCreateTables()`, `initGenerateData()`, partitioning, foreign keys, vacuuming | `make check` (TAP: 681 tests pass) |
| **Phase 7** | `engine.c` / `engine.h` & `pgbench.c` | `advanceConnectionState()`, `threadRun()`, async client step loop; `pgbench.c` reduced to clean CLI driver (~400 lines) | Full TAP + Regression verification |

---

## 5. Benefits for Future Workloads (TPC-C, TPC-E, TPC-H)

1. **Effortless Addition of New Meta-Commands**: Adding `\for`, `\while`, `\try`, `\catch` requires only adding one `CommandHandler` function and an entry in `commands.c`.
2. **Heterogeneous Workload Groups (`--group`)**: Multiple client groups executing independent script streams can run on top of clean `EngineContext` / `TState` instances without global state conflicts.
3. **Driver / Stored-Procedure Benchmarks**: Client state machine is decoupled from SQL execution, making it trivial to plug in batch call drivers or stored procedure executors.
4. **Enhanced Diagnostic & Profiling Tools**: Clean `stats.c` and `poller.c` abstractions make integrating monotonic timing, micro-benchmarking, and custom Prometheus/Monarch metrics straightforward.