memtier_benchmark.h

/*
 * Copyright (C) 2011-2026 Redis Labs Ltd.
 *
 * This file is part of memtier_benchmark.
 *
 * memtier_benchmark is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 2.
 *
 * memtier_benchmark is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with memtier_benchmark.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef _MEMTIER_BENCHMARK_H
#define _MEMTIER_BENCHMARK_H

#include <atomic>
#include <vector>
#include <string>
#include <utility>
#include <sys/time.h>
#include <pthread.h>
#include "config_types.h"

#ifdef USE_TLS
#include <openssl/ssl.h>
#endif

// Forward declaration
class statsd_client;
class prometheus_exporter;

#define LOGLEVEL_ERROR 0
#define LOGLEVEL_DEBUG 1

#define benchmark_debug_log(...) benchmark_log_file_line(LOGLEVEL_DEBUG, __FILE__, __LINE__, __VA_ARGS__)

#define benchmark_error_log(...) benchmark_log(LOGLEVEL_ERROR, __VA_ARGS__)

enum key_pattern_index
{
    key_pattern_set = 0,
    key_pattern_delimiter = 1,
    key_pattern_get = 2
};

enum PROTOCOL_TYPE
{
    PROTOCOL_REDIS_DEFAULT,
    PROTOCOL_RESP2,
    PROTOCOL_RESP3,
    PROTOCOL_MEMCACHE_TEXT,
    PROTOCOL_MEMCACHE_BINARY,
};

// ---------------------------------------------------------------------------
// Read-preference enums (storage lives in benchmark_config below)
// ---------------------------------------------------------------------------

// Which node class should receive read commands.
enum read_pref_mode
{
    rp_primary = 0,         // all reads go to the master/primary (default)
    rp_secondary,           // reads go to replica nodes only
    rp_secondary_preferred, // replicas preferred; fall back to primary
    rp_nearest              // any node; no strict placement guarantee
};

// What to do when the target node class is unavailable.
enum read_pref_fallback
{
    rpf_error = 0, // return an error to the caller (default)
    rpf_queue,     // queue the request until a suitable node is available
    rpf_primary    // fall back silently to the primary
};

// Lightweight representation of a --read-server HOST:PORT entry.
// Stored at parse time; routing code converts to server_addr objects later.
struct read_server_spec
{
    std::string host;
    unsigned short port;

    read_server_spec(const std::string &h, unsigned short p) : host(h), port(p) {}
};

// Shared MGET slot cache: built once (lazily, on first topology load) and
// read concurrently by all cluster_client threads.  m_mget_slot_keys is
// identical for every thread — only the per-slot round-robin cursors differ.
struct mget_slot_cache
{
    std::vector<std::vector<unsigned long long> > slot_keys; // [slot] → key indices; read-only after built
    std::atomic<bool> built;
    pthread_mutex_t mutex;

    mget_slot_cache() : built(false) { pthread_mutex_init(&mutex, NULL); }
    ~mget_slot_cache() { pthread_mutex_destroy(&mutex); }

private:
    mget_slot_cache(const mget_slot_cache &);
    mget_slot_cache &operator=(const mget_slot_cache &);
};

struct benchmark_config
{
    const char *server;
    unsigned short port;
    struct server_addr *server_addr;
    const char *unix_socket;
    int resolution;
    enum PROTOCOL_TYPE protocol;
    const char *out_file;
    const char *client_stats;
    unsigned int run_count;
    int debug;
    int show_config;
    int hide_histogram;
    config_quantiles print_percentiles;
    bool print_all_runs;
    bool realtime_latencies;
    int distinct_client_seed;
    int randomize;
    int next_client_idx;
    unsigned long long requests;
    unsigned int clients;
    unsigned int threads;
    unsigned int test_time;
    config_ratio ratio;
    unsigned int pipeline;
    unsigned int data_size;
    unsigned int data_offset;
    bool random_data;
    struct config_range data_size_range;
    config_weight_list data_size_list;
    const char *data_size_pattern;
    struct config_range expiry_range;
    const char *data_import;
    int data_verify;
    int verify_only;
    int generate_keys;
    const char *key_prefix;
    unsigned long long key_minimum;
    unsigned long long key_maximum;
    double key_stddev;
    double key_median;
    double key_zipf_exp;
    const char *key_pattern;
    unsigned int reconnect_interval;
    bool reconnect_on_error;
    unsigned int max_reconnect_attempts;
    double reconnect_backoff_factor;
    // Per-command retry (independent of reconnect_on_error):
    //   retry_on_error      master switch (default: off)
    //   max_retries         -1 = unlimited (default), 0 = disabled even with switch on, N>0 = bounded
    //   retry_backoff_ms    delay between retries, in milliseconds (default 0, immediate)
    //   retry_backoff_factor   exponential multiplier on retry_backoff_ms (default 0.0 = constant)
    //   retry_on_filter     NULL = built-in classifier (retry everything except permanent set);
    //                       non-NULL = comma-list of error-status prefixes to restrict retries to
    //   max_retry_queue     hard cap on per-connection retry queue (0 = pipeline * 4 default)
    bool retry_on_error;
    int max_retries;
    unsigned int retry_backoff_ms;
    double retry_backoff_factor;
    const char *retry_on_filter;
    unsigned int max_retry_queue;
    // When non-NULL, every request that ultimately fails (max_retries exhausted,
    // or permanent error like WRONGTYPE) is appended as a line of CSV to this
    // file. Off by default. Robust on errors: a failure to open or write is
    // logged once and the benchmark continues.
    const char *failed_keys_file;
    unsigned int connection_timeout;
    // Per-process bound on time spent in the *connection-setup* phase
    // (AUTH, HELLO, SELECT, CLUSTER SLOTS, initial probe). Once any thread
    // reaches steady-state (first non-setup response processed), this bound
    // no longer applies; --test-time takes over. Independent of
    // --connection-timeout (which is per-connect attempt) and --test-time
    // (which bounds the steady-state run). Default 30 s; 0 disables.
    unsigned int connection_stage_timeout;
    unsigned int thread_conn_start_min_jitter_micros;
    unsigned int thread_conn_start_max_jitter_micros;
    int multi_key_get;
    struct mget_slot_cache *mget_cache; // NULL unless cluster_mode && multi_key_get > 0
    const char *authenticate;
    int select_db;
    const char *uri;
    bool no_expiry;
    bool resolve_on_connect;
    // WAIT related
    config_ratio wait_ratio;
    config_range num_slaves;
    config_range wait_timeout;
    // JSON additions
    const char *json_out_file;
    bool cluster_mode;
    // When set together with --cluster-mode, every full rotation of --command
    // entries (one logical transactional unit, e.g. WATCH/MULTI/.../EXEC) is
    // pinned to a single shard connection so that keyless commands stay on
    // the same connection as the keyed ones.
    bool transaction;
    struct arbitrary_command_list *arbitrary_commands;
    const char *monitor_input;
    struct monitor_command_list *monitor_commands;
    char monitor_pattern;
    bool command_stats_by_type; // true = aggregate by command type (default), false = per command line
    bool command_miss_tracking; // true = auto (track misses for known shapes), false = off
    double miss_rate_threshold; // warn when miss rate exceeds this fraction (default 0.01 = 1%)
    double cpu_warn_threshold;  // warn when a memtier thread's CPU exceeds this fraction of a core (default 0.95)
    const char *hdr_prefix;
    unsigned int request_rate;
    unsigned int request_per_interval;
    unsigned int request_interval_microsecond;
    // Client staircase ramp-up
    unsigned int clients_start;
    unsigned int clients_step;
    unsigned int step_duration;
    struct timeval benchmark_start_time;
    // Index of the currently-executing run (1-based), set at the top of
    // run_benchmark(). Read by cluster_client::handle_cluster_slots() so the
    // one-shot topology summary it prints can label itself "[RUN #N]" and emit
    // exactly once per run (runs are sequential, so a plain field is enough).
    unsigned int current_run_id;
    // StatsD metrics export
    const char *statsd_host;
    unsigned short statsd_port;
    const char *statsd_prefix;
    const char *statsd_run_label;
    unsigned short graphite_port;
    statsd_client *statsd;
    // Prometheus metrics export (PLAN.md v5 sections 3.1, 5). The four flag
    // members are guarded; the exporter pointer is unguarded (statsd idiom,
    // Decisions #25): NULL when disabled or compiled out.
#ifdef HAVE_EVHTTP
    int prometheus_port;              // sentinel -1 = unset; 0 = ephemeral; else fixed port
    const char *prometheus_bind_addr; // NULL until config_init_defaults applies 127.0.0.1
    std::vector<std::pair<std::string, std::string> > prometheus_run_labels; // raw; renderer escapes
    std::vector<double> prometheus_latency_buckets;                          // seconds; empty = default
#endif
    prometheus_exporter *prometheus; // unguarded ownership alias of g_prom_exporter
    // SCAN incremental cursor iteration
    bool scan_incremental_iteration;
    unsigned int scan_incremental_max_iterations;
    arbitrary_command *scan_continuation_command;
    // ---------------------------------------------------------------------------
    // Read-preference configuration
    //   read_preference         which node class receives reads (default: primary)
    //   read_preference_fallback what to do when the target is unavailable
    //   read_servers            replica endpoints for standalone mode
    //                           (--read-server HOST:PORT, repeatable)
    //   replica_clients         connections per replica per client thread;
    //                           0 = inherit --clients
    //   replicas_per_shard      cap on replicas used per shard; 0 = all
    // ---------------------------------------------------------------------------
    enum read_pref_mode read_preference;
    enum read_pref_fallback read_preference_fallback;
    std::vector<read_server_spec> read_servers;
    unsigned int replica_clients;
    unsigned int replicas_per_shard;
#ifdef USE_TLS
    bool tls;
    const char *tls_cert;
    const char *tls_key;
    const char *tls_cacert;
    bool tls_skip_verify;
    const char *tls_sni;
    int tls_protocols;
    SSL_CTX *openssl_ctx;
    // Negotiated TLS protocol/cipher, captured once on the first completed
    // handshake (static OpenSSL strings; NULL until then). Written under a
    // call_once on a worker thread, read on the main thread post-join.
    const char *tls_negotiated_version;
    const char *tls_negotiated_cipher;
#endif
};


extern void benchmark_log_file_line(int level, const char *filename, unsigned int line, const char *fmt, ...);
extern void benchmark_log(int level, const char *fmt, ...);
bool is_redis_protocol(enum PROTOCOL_TYPE type);

// ---------------------------------------------------------------------------
// Connection-stage supervisor (Phase 1 of #426)
// ---------------------------------------------------------------------------
//
// The worker threads (shard_connection / cluster_client) report two events
// here:
//   - report_connection_stage_failure(): a connection-setup step failed
//     (AUTH / HELLO / SELECT / CLUSTER SLOTS / -ERR / parse error during
//     initial probe). The first call in a streak stamps the wall-clock
//     start of the streak; subsequent calls only update the last-error
//     message.
//   - report_connection_stage_success(): the worker exited the conn-setup
//     phase and processed a real response. Clears the streak and arms a
//     latch (steady_state_reached) so the supervisor stops policing this
//     run for setup-stalls.
//
// The main thread polls connection_stage_should_abort() once per second
// from run_benchmark(). It returns true when either:
//   (a) a failure streak has been live for >= --connection-stage-timeout, or
//   (b) the run has been alive for >= --connection-stage-timeout without
//       any thread ever reaching steady state (covers the "stuck WAIT" /
//       "first request never returns" hangs in #426 #17).
//
// All state lives behind atomics + a small mutex protecting the
// last-error string so worker threads can call into it lock-free on the
// hot success path.
void connection_stage_supervisor_reset(void);
void report_connection_stage_failure(const char *err);
void report_connection_stage_success(void);
bool connection_stage_should_abort(unsigned int timeout_secs, std::string *out_last_err, unsigned int *out_elapsed);

#endif /* _MEMTIER_BENCHMARK_H */