// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

// Concurrency stress test on a POOLED mysql channel: many bthreads hammering
// ONE pooled Channel, re-running concurrently the self-checking work units
// that the two sibling integration files cover serially:
//
//   * brpc_mysql_txn_integration_unittest.cpp       (transaction scenarios)
//   * brpc_mysql_prepared_integration_unittest.cpp  (prepared-statement
//                                                     scenarios)
//
// Only the literal data values differ from the siblings so that cross-talk
// between concurrent workers is detectable by value.  A POOLED Channel must
// check out / return pooled sockets without races, must PIN one socket per
// MysqlTransaction, and must keep concurrent transactions / prepared
// statements isolated.
//
// WHAT IT CHECKS
// ConnectionType = POOLED, pool capped at FIVE connections via the gflag
// `max_connection_pool_size` (DEFINE_int32 max_connection_pool_size in
// src/brpc/socket.cpp:99).  FIVE is deliberate: with more workers than
// pooled sockets we exercise BOTH pooled-socket reuse AND the create-a-NEW-
// connection-under-load path concurrently, surfacing checkout/return races,
// transaction connection-affinity/pinning under contention, and fd_version ABA.
//
//   * ManyWorkersMixedScenarios:
//       16-32 bthreads, each looping ~50x, each iteration picks ONE of five
//       representative self-checking work units (3 txn + 2 prepared) and
//       asserts its OWN correct, independent result.  Per-worker scratch tables
//       / per-iteration ids keep row-count assertions exact under concurrency.
//
//   * TwoTransactionsHoldDifferentPinnedSockets (focused check a):
//       two transactions in parallel must hold DIFFERENT pinned SocketIds
//       (GetSocketId()) and must not see each other's rows.
//
//   * TransactionPlusPreparedInParallel (focused check b):
//       one transaction + one prepared statement in parallel, each returns its
//       own correct independent result; the prepared path must not disturb the
//       transaction's pinned connection.
//
// The bar: NO concurrency bug across all iterations -- no shared-socket
// corruption, no interleaved/wrong replies, no crash.
//
// HARNESS
// Reuses the gflag-driven, self-spawning-mysqld harness from the two sibling
// integration files (flags -mysql_use_running_server / -mysql_host / -port /
// -user / -password / -schema; MysqlAuthenticator-based pooled Channel).  When
// no mysqld is reachable every test GTEST_SKIP()s, so the file is CI-safe.

#include <gflags/gflags.h>
#include <gtest/gtest.h>

#include <arpa/inet.h>
#include <netinet/in.h>
#include <pthread.h>
#include <sys/socket.h>
#include <unistd.h>

#include <cstdio>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>

#include "brpc/channel.h"
#include "brpc/controller.h"
#include "brpc/policy/mysql/mysql.h"
#include "brpc/policy/mysql/mysql_transaction.h"
#include "brpc/policy/mysql/mysql_authenticator.h"
#include "bthread/bthread.h"
#include "butil/logging.h"
#include "butil/string_printf.h"
#include "butil/strings/string_piece.h"

// Flags mirror the sibling integration files so one command line drives them
// all against the same server.  Each *_unittest.cpp links into its own binary
// (the test/ CMake glob), so re-declaring these flags here is not a clash.
DEFINE_bool(mysql_use_running_server, false,
            "Use an already-running MySQL server instead of spawning a "
            "throwaway one; the running server is neither started nor stopped "
            "by the test.");
DEFINE_string(mysql_host, "127.0.0.1",
              "Host of the running MySQL server "
              "(only with -mysql_use_running_server).");
DEFINE_int32(mysql_port, 13306,
             "TCP port of the MySQL server (used for both the running server "
             "and the spawned throwaway server).");
DEFINE_string(mysql_user, "root", "Login user for the concurrency tests.");
DEFINE_string(mysql_password, "",
              "Password for -mysql_user (empty for the spawned server).");
DEFINE_string(mysql_schema, "brpc_pool_conc_test",
              "Schema (database) the concurrency tests create and use.");

namespace {

#ifndef GFLAGS_NS
#define GFLAGS_NS GFLAGS_NAMESPACE
#endif

#define MYSQLD_BIN "mysqld"

static const char* kCollation = "utf8mb4_general_ci";

// Concurrency knobs.  kWorkers is deliberately > the pool cap (5) so workers
// contend for pooled sockets AND force new-connection creation under load.
const int kWorkers = 24;
const int kIterationsPerWorker = 50;
const int kPoolCap = 5;

// Throwaway-server harness (mirrors the two sibling integration files, which
// mirror brpc_redis_unittest.cpp).  >0: forked pid; -2: external running
// server reachable; -1: no server -> tests skip.
static pthread_once_t g_start_once = PTHREAD_ONCE_INIT;
static pid_t g_mysqld_pid = -1;
static std::string g_host = "127.0.0.1";
static int g_port = 13306;
static std::string g_user = "root";
static std::string g_password;
static std::string g_schema;

static std::string TestDataDir() {
    char cwd[1024];
    if (getcwd(cwd, sizeof(cwd)) == NULL) {
        return std::string("/tmp/mysql_pool_conc_data_for_test");
    }
    return std::string(cwd) + "/mysql_pool_conc_data_for_test";
}

static void RemoveMysqlServer() {
    if (g_mysqld_pid > 0) {
        puts("[Stopping mysqld]");
        char cmd[1280];
        snprintf(cmd, sizeof(cmd), "kill %d", g_mysqld_pid);
        CHECK(0 == system(cmd));
        usleep(500000);
        snprintf(cmd, sizeof(cmd), "rm -rf '%s'", TestDataDir().c_str());
        CHECK(0 == system(cmd));
    }
}

// Raw TCP probe for server readiness; returns fd (caller closes) or -1.
static int ProbeConnect() {
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        return -1;
    }
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(static_cast<uint16_t>(g_port));
    addr.sin_addr.s_addr = inet_addr(g_host.c_str());
    if (connect(fd, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
        close(fd);
        return -1;
    }
    return fd;
}

static void StartServerOnce() {
    if (FLAGS_mysql_use_running_server) {
        g_host = FLAGS_mysql_host;
        g_port = FLAGS_mysql_port;
        g_user = FLAGS_mysql_user;
        g_password = FLAGS_mysql_password;
        g_schema = FLAGS_mysql_schema;
        printf("[Using running mysqld at %s:%d as user '%s', schema '%s']\n",
               g_host.c_str(), g_port, g_user.c_str(), g_schema.c_str());
        int fd = ProbeConnect();
        if (fd >= 0) {
            close(fd);
            g_mysqld_pid = -2;
        } else {
            printf("Cannot reach running mysqld at %s:%d, tests will skip\n",
                   g_host.c_str(), g_port);
        }
        return;
    }

    if (system("which " MYSQLD_BIN) != 0) {
        puts("Fail to find " MYSQLD_BIN ", concurrency tests will be skipped");
        return;
    }
    g_host = "127.0.0.1";
    g_port = FLAGS_mysql_port;
    g_user = "root";
    g_password.clear();
    g_schema = FLAGS_mysql_schema;
    const std::string datadir = TestDataDir();
    char cmd[2048];
    snprintf(cmd, sizeof(cmd), "rm -rf '%s' && mkdir -p '%s'",
             datadir.c_str(), datadir.c_str());
    if (system(cmd) != 0) {
        puts("Fail to create datadir, concurrency tests will be skipped");
        return;
    }
    snprintf(cmd, sizeof(cmd),
             MYSQLD_BIN " --initialize-insecure --datadir='%s'"
             " --log-error='%s/init.err'",
             datadir.c_str(), datadir.c_str());
    if (system(cmd) != 0) {
        puts("Fail to initialize mysqld datadir, tests will be skipped");
        snprintf(cmd, sizeof(cmd), "rm -rf '%s'", datadir.c_str());
        CHECK(0 == system(cmd));
        return;
    }
    atexit(RemoveMysqlServer);

    g_mysqld_pid = fork();
    if (g_mysqld_pid < 0) {
        puts("Fail to fork");
        exit(1);
    } else if (g_mysqld_pid == 0) {
        puts("[Starting mysqld]");
        char port_arg[32];
        snprintf(port_arg, sizeof(port_arg), "--port=%d", FLAGS_mysql_port);
        const std::string datadir_arg = "--datadir=" + datadir;
        const std::string socket_arg = "--socket=" + datadir + "/mysqld.sock";
        const std::string pidfile_arg = "--pid-file=" + datadir + "/mysqld.pid";
        const std::string logerr_arg = "--log-error=" + datadir + "/mysqld.err";
        char* const argv[] = {
            (char*)MYSQLD_BIN,
            (char*)datadir_arg.c_str(),
            (char*)port_arg,
            (char*)socket_arg.c_str(),
            (char*)pidfile_arg.c_str(),
            (char*)logerr_arg.c_str(),
            (char*)"--mysqlx=OFF",
            (char*)"--bind-address=127.0.0.1",
            NULL};
        if (execvp(MYSQLD_BIN, argv) < 0) {
            puts("Fail to run " MYSQLD_BIN);
            exit(1);
        }
    }
    for (int i = 0; i < 300; ++i) {
        int fd = ProbeConnect();
        if (fd >= 0) {
            close(fd);
            return;
        }
        usleep(100000);
    }
    puts("mysqld did not become ready, concurrency tests will be skipped");
    g_mysqld_pid = -1;
}

// Small helpers over the brpc MySQL public API.

// Plain (no transaction / no statement) query on a fresh pooled connection.
static bool RunPlain(brpc::Channel& channel, const std::string& sql,
                     brpc::MysqlResponse* resp, std::string* err) {
    brpc::MysqlRequest req;
    if (!req.Query(sql)) {
        if (err) *err = "build query failed: " + sql;
        return false;
    }
    brpc::Controller cntl;
    channel.CallMethod(NULL, &cntl, &req, resp, NULL);
    if (cntl.Failed()) {
        if (err) *err = "rpc failed: " + cntl.ErrorText();
        return false;
    }
    if (resp->reply_size() < 1) {
        if (err) *err = "no reply for: " + sql;
        return false;
    }
    if (resp->reply(0).is_error()) {
        const brpc::MysqlReply& r = resp->reply(0);
        if (err) {
            *err = butil::string_printf("mysql error %u: %.*s (sql=%s)",
                r.error().errcode(), (int)r.error().msg().size(),
                r.error().msg().data(), sql.c_str());
        }
        return false;
    }
    return true;
}

// |sql| INSIDE transaction |tx| (its pinned connection).
static bool RunInTx(brpc::Channel& channel, const brpc::MysqlTransaction* tx,
                    const std::string& sql, brpc::MysqlResponse* resp,
                    std::string* err) {
    brpc::MysqlRequest req(tx);
    if (!req.Query(sql)) {
        if (err) *err = "build query failed: " + sql;
        return false;
    }
    brpc::Controller cntl;
    channel.CallMethod(NULL, &cntl, &req, resp, NULL);
    if (cntl.Failed()) {
        if (err) *err = "rpc failed: " + cntl.ErrorText();
        return false;
    }
    if (resp->reply_size() < 1) {
        if (err) *err = "no reply for: " + sql;
        return false;
    }
    return true;
}

// Row count of the first reply when it is a result set, else -1.
static int64_t ResultRowCount(const brpc::MysqlResponse& resp) {
    if (resp.reply_size() < 1) {
        return -1;
    }
    const brpc::MysqlReply& r = resp.reply(0);
    if (!r.is_resultset()) {
        return -1;
    }
    return static_cast<int64_t>(r.row_count());
}

// Coerce an integer-ish field to long long (handles the various widths the
// server may choose for a column / expression).
static bool FieldToLongLong(const brpc::MysqlReply::Field& f, long long* out) {
    if (f.is_sbigint())       *out = f.sbigint();
    else if (f.is_bigint())   *out = (long long)f.bigint();
    else if (f.is_sinteger()) *out = f.sinteger();
    else if (f.is_integer())  *out = (long long)f.integer();
    else if (f.is_string())   *out = atoll(f.string().as_string().c_str());
    else return false;
    return true;
}

static int InitPooledChannel(brpc::Channel* channel,
                             brpc::policy::MysqlAuthenticator** out_auth,
                             const std::string& schema) {
    brpc::policy::MysqlAuthenticator* auth =
        new brpc::policy::MysqlAuthenticator(g_user, g_password, schema, "",
                                             kCollation);
    *out_auth = auth;
    brpc::ChannelOptions options;
    options.protocol = brpc::PROTOCOL_MYSQL;
    options.connection_type = "pooled";  // CONNECTION_TYPE_POOLED
    options.auth = auth;
    options.timeout_ms = 10000;
    options.connect_timeout_ms = 5000;
    options.max_retry = 0;
    return channel->Init(g_host.c_str(), g_port, &options);
}

// Fixture: one shared pooled channel, pool capped at FIVE.  Per-test scratch
// tables are created by the tests themselves (per-worker, to keep row counts
// exact under concurrency).
class MysqlPoolConcurrencyTest : public testing::Test {
protected:
    static bool NoServer() { return g_mysqld_pid == -1; }

    void SetUp() override {
        pthread_once(&g_start_once, StartServerOnce);
        if (NoServer()) {
            GTEST_SKIP() << "no mysqld available; skipping pool-concurrency "
                            "integration test (set -mysql_use_running_server "
                            "or install mysqld)";
        }
        // Cap the pool at FIVE for the whole test.  Verified flag name:
        // src/brpc/socket.cpp:99 DEFINE_int32(max_connection_pool_size, ...).
        ASSERT_FALSE(GFLAGS_NS::SetCommandLineOption(
                         "max_connection_pool_size",
                         std::to_string(kPoolCap).c_str()).empty())
            << "failed to set gflag max_connection_pool_size";

        // Create the schema over a schema-less channel, then bind to it.
        brpc::policy::MysqlAuthenticator* setup_auth = NULL;
        ASSERT_EQ(0, InitPooledChannel(&_setup_channel, &setup_auth, ""));
        _setup_auth.reset(setup_auth);
        brpc::MysqlResponse resp;
        std::string err;
        ASSERT_TRUE(RunPlain(_setup_channel,
                             "CREATE DATABASE IF NOT EXISTS " + g_schema,
                             &resp, &err)) << err;

        brpc::policy::MysqlAuthenticator* auth = NULL;
        ASSERT_EQ(0, InitPooledChannel(&_channel, &auth, g_schema));
        _auth.reset(auth);
    }

    brpc::Channel _setup_channel;
    brpc::Channel _channel;
    // Authenticators must outlive the channels that point at them.
    std::unique_ptr<brpc::policy::MysqlAuthenticator> _setup_auth;
    std::unique_ptr<brpc::policy::MysqlAuthenticator> _auth;
};

// WORK UNITS
// Each work unit is a self-checking re-run of a sibling scenario, with its OWN
// independent expected result so cross-talk/corruption is detectable.  Each
// returns true on a correct result; on any failure it fills |err|.
//
// All work units use a PER-WORKER scratch table (passed in) so concurrent
// workers never share rows, keeping row-count assertions exact.

// WU1 -- txn commit makes rows visible.
// (Transaction-commit-visibility check; uses its own per-worker id 71xxx
//  'aria' so concurrent workers never collide on data.)
static bool WU_TxnCommitVisible(brpc::Channel& ch, const std::string& table,
                                int iter, std::string* err) {
    const int id = 71000 + iter;
    const char* name = "aria";
    brpc::MysqlResponse resp;
    brpc::MysqlTransactionUniquePtr tx =
        brpc::NewMysqlTransaction(ch, brpc::MysqlTransactionOptions());
    if (tx == NULL) { *err = "WU1: NewMysqlTransaction NULL"; return false; }
    if (!RunInTx(ch, tx.get(),
                 butil::string_printf("INSERT INTO %s VALUES (%d, '%s')",
                                      table.c_str(), id, name),
                 &resp, err)) return false;
    if (resp.reply(0).is_error()) {
        *err = "WU1 INSERT err: " + resp.reply(0).error().msg().as_string();
        return false;
    }
    if (!tx->commit()) { *err = "WU1: commit failed"; return false; }

    // A fresh pooled connection must now see exactly our committed row.
    if (!RunPlain(ch, butil::string_printf(
                          "SELECT id, name FROM %s WHERE id=%d",
                          table.c_str(), id), &resp, err)) return false;
    if (ResultRowCount(resp) != 1) {
        *err = butil::string_printf("WU1: expected 1 visible row, got %lld",
                                    (long long)ResultRowCount(resp));
        return false;
    }
    const brpc::MysqlReply::Row& row = resp.reply(0).next();
    long long got_id = 0;
    if (!FieldToLongLong(row.field(0), &got_id) || got_id != id) {
        *err = "WU1: wrong id read back"; return false;
    }
    if (row.field(1).string().as_string() != name) {
        *err = "WU1: wrong name read back"; return false;
    }
    return true;
}

// WU2 -- txn rollback discards the insert.
// (Rollback-discards-insert check; uses per-worker id 82xxx 'cory'.)
static bool WU_TxnRollbackDiscards(brpc::Channel& ch, const std::string& table,
                                   int iter, std::string* err) {
    const int id = 82000 + iter;
    brpc::MysqlResponse resp;
    brpc::MysqlTransactionUniquePtr tx =
        brpc::NewMysqlTransaction(ch, brpc::MysqlTransactionOptions());
    if (tx == NULL) { *err = "WU2: NewMysqlTransaction NULL"; return false; }
    if (!RunInTx(ch, tx.get(),
                 butil::string_printf("INSERT INTO %s VALUES (%d, 'cory')",
                                      table.c_str(), id),
                 &resp, err)) return false;
    if (resp.reply(0).is_error()) {
        *err = "WU2 INSERT err: " + resp.reply(0).error().msg().as_string();
        return false;
    }
    if (!tx->rollback()) { *err = "WU2: rollback failed"; return false; }

    // The rolled-back row must be gone on a fresh connection.
    if (!RunPlain(ch, butil::string_printf("SELECT id FROM %s WHERE id=%d",
                                           table.c_str(), id),
                  &resp, err)) return false;
    if (ResultRowCount(resp) != 0) {
        *err = butil::string_printf(
            "WU2: rolled-back insert still visible (%lld rows)",
            (long long)ResultRowCount(resp));
        return false;
    }
    return true;
}

// WU3 -- a SELECT inside an open txn sees the txn's own uncommitted write.
// (Read-your-own-write check; uses per-worker id 93xxx 'echo'.)
static bool WU_TxnReadsOwnWrite(brpc::Channel& ch, const std::string& table,
                                int iter, std::string* err) {
    const int id = 93000 + iter;
    const char* name = "echo";
    brpc::MysqlResponse resp;
    brpc::MysqlTransactionUniquePtr tx =
        brpc::NewMysqlTransaction(ch, brpc::MysqlTransactionOptions());
    if (tx == NULL) { *err = "WU3: NewMysqlTransaction NULL"; return false; }
    if (!RunInTx(ch, tx.get(),
                 butil::string_printf("INSERT INTO %s VALUES (%d, '%s')",
                                      table.c_str(), id, name),
                 &resp, err)) return false;
    if (resp.reply(0).is_error()) {
        *err = "WU3 INSERT err: " + resp.reply(0).error().msg().as_string();
        return false;
    }
    // Same pinned connection: must see its own uncommitted row.
    if (!RunInTx(ch, tx.get(),
                 butil::string_printf("SELECT name FROM %s WHERE id=%d",
                                      table.c_str(), id),
                 &resp, err)) return false;
    bool ok = (ResultRowCount(resp) == 1) &&
              resp.reply(0).next().field(0).string().as_string() == name;
    // Discard the row so the per-worker table stays empty for the next pick.
    tx->rollback();
    if (!ok) { *err = "WU3: txn did not read its own uncommitted write"; }
    return ok;
}

// WU4 -- prepared: bind one INT param, fetch the matching row.
// (Prepared-bind-int check; each worker seeds its OWN (id 64xxx, 'lyra') row
//  and binds it back so the result is unique per worker.)
static bool WU_PreparedBindInt(brpc::Channel& ch, const std::string& table,
                               int iter, std::string* err) {
    const int id = 64000 + iter;
    const char* name = "lyra";
    brpc::MysqlResponse resp;
    // Seed our own row (autocommit) then read it back via a prepared SELECT.
    if (!RunPlain(ch, butil::string_printf("INSERT INTO %s VALUES (%d, '%s')",
                                           table.c_str(), id, name),
                  &resp, err)) return false;

    brpc::MysqlStatementUniquePtr stmt = brpc::NewMysqlStatement(
        ch, butil::string_printf("SELECT name FROM %s WHERE id=?",
                                 table.c_str()));
    if (stmt == NULL) { *err = "WU4: NewMysqlStatement NULL"; return false; }
    if (stmt->param_count() != 1u) { *err = "WU4: param_count != 1"; return false; }

    brpc::MysqlRequest req(stmt.get());
    if (!req.AddParam((int32_t)id)) { *err = "WU4: AddParam failed"; return false; }
    brpc::Controller cntl;
    ch.CallMethod(NULL, &cntl, &req, &resp, NULL);
    if (cntl.Failed()) { *err = "WU4 rpc: " + cntl.ErrorText(); return false; }
    if (resp.reply_size() < 1 || !resp.reply(0).is_resultset()) {
        *err = "WU4: not a resultset"; return false;
    }
    if (resp.reply(0).row_count() != 1u) {
        *err = "WU4: expected exactly 1 row"; return false;
    }
    bool ok = resp.reply(0).next().field(0).string().as_string() == name;
    // Clean our seeded row so subsequent picks on this table stay exact.
    RunPlain(ch, butil::string_printf("DELETE FROM %s WHERE id=%d",
                                      table.c_str(), id), &resp, err);
    if (!ok) { *err = "WU4: prepared bind returned wrong/cross-talked name"; }
    return ok;
}

// WU5 -- prepared multi-param arithmetic: SELECT ? + ? with two INT params.
// (Two-param arithmetic check; uses per-iteration operands so each worker
//  verifies its OWN sum.)
static bool WU_PreparedArithmetic(brpc::Channel& ch, int worker, int iter,
                                  std::string* err) {
    const int32_t a = 1000 + worker;
    const int32_t b = 7 + iter;
    const long long expect = (long long)a + b;
    brpc::MysqlStatementUniquePtr stmt =
        brpc::NewMysqlStatement(ch, "SELECT CAST(? AS SIGNED) + CAST(? AS SIGNED)");
    if (stmt == NULL) { *err = "WU5: NewMysqlStatement NULL"; return false; }
    if (stmt->param_count() != 2u) { *err = "WU5: param_count != 2"; return false; }

    brpc::MysqlRequest req(stmt.get());
    if (!req.AddParam(a) || !req.AddParam(b)) {
        *err = "WU5: AddParam failed"; return false;
    }
    brpc::MysqlResponse resp;
    brpc::Controller cntl;
    ch.CallMethod(NULL, &cntl, &req, &resp, NULL);
    if (cntl.Failed()) { *err = "WU5 rpc: " + cntl.ErrorText(); return false; }
    if (resp.reply_size() < 1 || !resp.reply(0).is_resultset() ||
        resp.reply(0).row_count() != 1u) {
        *err = "WU5: bad resultset"; return false;
    }
    long long got = 0;
    if (!FieldToLongLong(resp.reply(0).next().field(0), &got) || got != expect) {
        *err = butil::string_printf("WU5: ?+? wrong (got %lld want %lld)",
                                    got, expect);
        return false;
    }
    return true;
}

// Worker driver for ManyWorkersMixedScenarios.
struct MixWorkerArgs {
    brpc::Channel* channel;
    int worker_id;
    std::string table;   // per-worker scratch table
    std::string error;   // first failure (empty == all good)
    int completed;       // iterations completed without error
};

void* MixWorker(void* p) {
    MixWorkerArgs* a = static_cast<MixWorkerArgs*>(p);
    a->error.clear();
    a->completed = 0;
    for (int iter = 0; iter < kIterationsPerWorker; ++iter) {
        // Rotate through the five work units; offset by worker so different
        // workers run different units at the same instant.
        const int pick = (a->worker_id + iter) % 5;
        std::string err;
        bool ok = false;
        switch (pick) {
            case 0: ok = WU_TxnCommitVisible(*a->channel, a->table, iter, &err); break;
            case 1: ok = WU_TxnRollbackDiscards(*a->channel, a->table, iter, &err); break;
            case 2: ok = WU_TxnReadsOwnWrite(*a->channel, a->table, iter, &err); break;
            case 3: ok = WU_PreparedBindInt(*a->channel, a->table, iter, &err); break;
            default: ok = WU_PreparedArithmetic(*a->channel, a->worker_id, iter, &err); break;
        }
        if (!ok) {
            a->error = butil::string_printf("worker %d iter %d pick %d: %s",
                                            a->worker_id, iter, pick, err.c_str());
            return NULL;
        }
        ++a->completed;
    }
    return NULL;
}

// TEST 1: many bthreads, each looping ~50x over a mix of the reused work
// units, on ONE pooled channel capped at 5 sockets.  Surfaces checkout/return
// races, affinity-under-contention bugs, fd_version ABA, and the new-connection
// creation path.
TEST_F(MysqlPoolConcurrencyTest, ManyWorkersMixedScenarios) {
    std::string err;
    std::vector<MixWorkerArgs> args(kWorkers);
    // One scratch table per worker so concurrent workers never share rows.
    for (int w = 0; w < kWorkers; ++w) {
        args[w].channel = &_channel;
        args[w].worker_id = w;
        args[w].table = butil::string_printf("pool_conc_w%d", w);
        brpc::MysqlResponse resp;
        ASSERT_TRUE(RunPlain(_channel, "DROP TABLE IF EXISTS " + args[w].table,
                             &resp, &err)) << err;
        ASSERT_TRUE(RunPlain(_channel,
                             "CREATE TABLE " + args[w].table +
                                 " (id INT PRIMARY KEY, name VARCHAR(32)) "
                                 "ENGINE=InnoDB",
                             &resp, &err)) << err;
    }

    std::vector<bthread_t> threads(kWorkers);
    for (int w = 0; w < kWorkers; ++w) {
        ASSERT_EQ(0, bthread_start_background(&threads[w], NULL, MixWorker,
                                              &args[w]));
    }
    for (int w = 0; w < kWorkers; ++w) {
        bthread_join(threads[w], NULL);
    }

    for (int w = 0; w < kWorkers; ++w) {
        EXPECT_TRUE(args[w].error.empty()) << args[w].error;
        EXPECT_EQ(kIterationsPerWorker, args[w].completed)
            << "worker " << w << " did not finish all iterations";
    }

    // Cleanup.
    for (int w = 0; w < kWorkers; ++w) {
        brpc::MysqlResponse resp;
        RunPlain(_channel, "DROP TABLE IF EXISTS " + args[w].table, &resp, &err);
    }
}

// Focused-check worker A: one transaction, per-worker table, INSERT + SELECT,
// records its pinned SocketId and the value it read.
struct AffinityWorkerArgs {
    brpc::Channel* channel;
    std::string table;
    int id;                    // value inserted (and expected back)
    brpc::SocketId socket_id;  // pinned socket for this txn
    int64_t row_count;
    int read_value;
    bool committed;
    std::string error;
};

void* AffinityWorker(void* p) {
    AffinityWorkerArgs* a = static_cast<AffinityWorkerArgs*>(p);
    a->error.clear();
    a->socket_id = 0;
    a->row_count = -1;
    a->read_value = -1;
    a->committed = false;

    brpc::MysqlTransactionUniquePtr tx =
        brpc::NewMysqlTransaction(*a->channel, brpc::MysqlTransactionOptions());
    if (tx == NULL) { a->error = "NewMysqlTransaction NULL"; return NULL; }
    a->socket_id = tx->GetSocketId();

    brpc::MysqlResponse resp;
    if (!RunInTx(*a->channel, tx.get(),
                 butil::string_printf("INSERT INTO %s VALUES (%d)",
                                      a->table.c_str(), a->id),
                 &resp, &a->error)) return NULL;
    if (resp.reply(0).is_error()) {
        a->error = "INSERT err: " + resp.reply(0).error().msg().as_string();
        return NULL;
    }
    // Read inside the txn: must see exactly our own row (per-worker table).
    if (!RunInTx(*a->channel, tx.get(),
                 butil::string_printf("SELECT v FROM %s", a->table.c_str()),
                 &resp, &a->error)) return NULL;
    a->row_count = ResultRowCount(resp);
    if (a->row_count == 1) {
        long long v = 0;
        if (FieldToLongLong(resp.reply(0).next().field(0), &v)) {
            a->read_value = (int)v;
        }
    }
    a->committed = tx->commit();
    if (!a->committed) a->error = "commit failed";
    return NULL;
}

// TEST 2 (focused check a): two transactions in parallel must hold DIFFERENT
// pinned SocketIds and must not see each other's rows.
TEST_F(MysqlPoolConcurrencyTest, TwoTransactionsHoldDifferentPinnedSockets) {
    const std::string t0 = "pool_conc_affinity_a";
    const std::string t1 = "pool_conc_affinity_b";
    std::string err;
    brpc::MysqlResponse resp;
    for (const std::string& t : {t0, t1}) {
        ASSERT_TRUE(RunPlain(_channel, "DROP TABLE IF EXISTS " + t, &resp, &err)) << err;
        ASSERT_TRUE(RunPlain(_channel,
                             "CREATE TABLE " + t + " (v INT) ENGINE=InnoDB",
                             &resp, &err)) << err;
    }

    for (int iter = 0; iter < kIterationsPerWorker; ++iter) {
        ASSERT_TRUE(RunPlain(_channel, "TRUNCATE TABLE " + t0, &resp, &err)) << err;
        ASSERT_TRUE(RunPlain(_channel, "TRUNCATE TABLE " + t1, &resp, &err)) << err;

        AffinityWorkerArgs a0{&_channel, t0, 30100 + iter, 0, -1, -1, false, ""};
        AffinityWorkerArgs a1{&_channel, t1, 30900 + iter, 0, -1, -1, false, ""};

        bthread_t b0, b1;
        ASSERT_EQ(0, bthread_start_background(&b0, NULL, AffinityWorker, &a0));
        ASSERT_EQ(0, bthread_start_background(&b1, NULL, AffinityWorker, &a1));
        bthread_join(b0, NULL);
        bthread_join(b1, NULL);

        ASSERT_TRUE(a0.error.empty()) << "iter " << iter << " txn0: " << a0.error;
        ASSERT_TRUE(a1.error.empty()) << "iter " << iter << " txn1: " << a1.error;

        // Each saw exactly its own single row -- no cross-talk.
        EXPECT_EQ(1, a0.row_count) << "iter " << iter;
        EXPECT_EQ(1, a1.row_count) << "iter " << iter;
        EXPECT_EQ(a0.id, a0.read_value) << "iter " << iter;
        EXPECT_EQ(a1.id, a1.read_value) << "iter " << iter;
        EXPECT_TRUE(a0.committed) << "iter " << iter;
        EXPECT_TRUE(a1.committed) << "iter " << iter;

        // Connection affinity: two concurrent txns hold DIFFERENT pinned sockets.
        EXPECT_NE(0u, a0.socket_id) << "iter " << iter;
        EXPECT_NE(0u, a1.socket_id) << "iter " << iter;
        EXPECT_NE(a0.socket_id, a1.socket_id)
            << "iter " << iter
            << ": two concurrent transactions shared a pooled socket! sid0="
            << a0.socket_id << " sid1=" << a1.socket_id;
    }

    for (const std::string& t : {t0, t1}) {
        RunPlain(_channel, "DROP TABLE IF EXISTS " + t, &resp, &err);
    }
}

// Focused-check worker B: a prepared statement (SELECT ? + ?) run repeatedly
// in parallel with a transaction; must return its own correct sum each time.
struct PreparedWorkerArgs {
    brpc::Channel* channel;
    int base;              // operand seed
    bool ok;
    std::string error;
};

void* PreparedWorker(void* p) {
    PreparedWorkerArgs* a = static_cast<PreparedWorkerArgs*>(p);
    a->ok = true;
    a->error.clear();
    for (int k = 0; k < 4; ++k) {
        std::string err;
        if (!WU_PreparedArithmetic(*a->channel, a->base, k, &err)) {
            a->ok = false;
            a->error = err;
            return NULL;
        }
    }
    return NULL;
}

// TEST 3 (focused check b): one transaction + one prepared statement in
// parallel each return correct independent results; the prepared path must not
// disturb the transaction's pinned connection.
TEST_F(MysqlPoolConcurrencyTest, TransactionPlusPreparedInParallel) {
    const std::string t = "pool_conc_txn_stmt";
    std::string err;
    brpc::MysqlResponse resp;
    ASSERT_TRUE(RunPlain(_channel, "DROP TABLE IF EXISTS " + t, &resp, &err)) << err;
    ASSERT_TRUE(RunPlain(_channel, "CREATE TABLE " + t + " (v INT) ENGINE=InnoDB",
                         &resp, &err)) << err;

    for (int iter = 0; iter < kIterationsPerWorker; ++iter) {
        ASSERT_TRUE(RunPlain(_channel, "TRUNCATE TABLE " + t, &resp, &err)) << err;

        AffinityWorkerArgs ta{&_channel, t, 50500 + iter, 0, -1, -1, false, ""};
        PreparedWorkerArgs pa{&_channel, 200 + iter, true, ""};

        bthread_t bt, bp;
        ASSERT_EQ(0, bthread_start_background(&bt, NULL, AffinityWorker, &ta));
        ASSERT_EQ(0, bthread_start_background(&bp, NULL, PreparedWorker, &pa));
        bthread_join(bt, NULL);
        bthread_join(bp, NULL);

        ASSERT_TRUE(ta.error.empty()) << "iter " << iter << " txn: " << ta.error;
        ASSERT_TRUE(pa.ok) << "iter " << iter << " prepared: " << pa.error;

        // Transaction result intact and isolated.
        EXPECT_EQ(1, ta.row_count) << "iter " << iter;
        EXPECT_EQ(ta.id, ta.read_value) << "iter " << iter;
        EXPECT_TRUE(ta.committed) << "iter " << iter;
        EXPECT_NE(0u, ta.socket_id) << "iter " << iter;
    }

    RunPlain(_channel, "DROP TABLE IF EXISTS " + t, &resp, &err);
}

// OWNER-PRIORITY CONCURRENCY CHECKS (TEST A / B / C)
//
// These three tests intentionally cap the pool SMALL relative to the number of
// concurrent workers (TEST A/B: 4 sockets vs 8 workers; TEST C: 2 sockets) so
// that workers both CONTEND for pooled sockets and FORCE new-connection
// creation, while transactions RESERVE (pull out) pooled sockets.  They use
// their OWN data namespace (ids in the 120xxx/130xxx/140xxx ranges, names
// "nova"/"zephyr"/"quill") that is distinct from the reused work units above
// (71xxx/82xxx/93xxx/64xxx; aria/cory/echo/lyra) and from the sibling files.
//
// Pool size is set PER-TEST via SetCommandLineOption at the top of each test
// body (tests may share a process, so the previous test's value must not leak
// in).  No ASSERT_* runs inside a bthread; every worker records into its args
// struct and the main thread asserts after join.

static void SetPoolCap(int cap) {
    GFLAGS_NS::SetCommandLineOption("max_connection_pool_size",
                                    std::to_string(cap).c_str());
}

// TEST A worker: one transaction running SEVERAL statements into its OWN
// per-worker scratch table, recording the pinned SocketId seen BEFORE every
// statement so the main thread can check intra-txn pinning and inter-txn
// isolation.
struct PinnedTxnWorkerArgs {
    brpc::Channel* channel;
    std::string table;             // per-worker scratch table
    int id;                        // per-worker row id (unique)
    std::vector<brpc::SocketId> socket_ids;  // GetSocketId() before each stmt
    std::string error;             // empty == ok
};

void* PinnedTxnWorker(void* p) {
    PinnedTxnWorkerArgs* a = static_cast<PinnedTxnWorkerArgs*>(p);
    a->error.clear();
    a->socket_ids.clear();

    brpc::MysqlTransactionUniquePtr tx =
        brpc::NewMysqlTransaction(*a->channel, brpc::MysqlTransactionOptions());
    if (tx == NULL) { a->error = "NewMysqlTransaction NULL"; return NULL; }

    brpc::MysqlResponse resp;

    // Statement 1: INSERT our own row.
    a->socket_ids.push_back(tx->GetSocketId());
    if (!RunInTx(*a->channel, tx.get(),
                 butil::string_printf("INSERT INTO %s VALUES (%d, 'nova')",
                                      a->table.c_str(), a->id),
                 &resp, &a->error)) return NULL;
    if (resp.reply(0).is_error()) {
        a->error = "INSERT err: " + resp.reply(0).error().msg().as_string();
        return NULL;
    }

    // Statement 2: SELECT our own (uncommitted) row back.
    a->socket_ids.push_back(tx->GetSocketId());
    if (!RunInTx(*a->channel, tx.get(),
                 butil::string_printf("SELECT name FROM %s WHERE id=%d",
                                      a->table.c_str(), a->id),
                 &resp, &a->error)) return NULL;
    if (ResultRowCount(resp) != 1 ||
        resp.reply(0).next().field(0).string().as_string() != "nova") {
        a->error = "SELECT-own-row did not read back its own write";
        return NULL;
    }

    // Statement 3: UPDATE our own row.
    a->socket_ids.push_back(tx->GetSocketId());
    if (!RunInTx(*a->channel, tx.get(),
                 butil::string_printf("UPDATE %s SET name='zephyr' WHERE id=%d",
                                      a->table.c_str(), a->id),
                 &resp, &a->error)) return NULL;
    if (resp.reply(0).is_error()) {
        a->error = "UPDATE err: " + resp.reply(0).error().msg().as_string();
        return NULL;
    }

    // Statement 4: SELECT the updated value back.
    a->socket_ids.push_back(tx->GetSocketId());
    if (!RunInTx(*a->channel, tx.get(),
                 butil::string_printf("SELECT name FROM %s WHERE id=%d",
                                      a->table.c_str(), a->id),
                 &resp, &a->error)) return NULL;
    if (ResultRowCount(resp) != 1 ||
        resp.reply(0).next().field(0).string().as_string() != "zephyr") {
        a->error = "SELECT after UPDATE saw wrong value";
        return NULL;
    }

    // Discard so the per-worker table is empty for the next outer-loop pass.
    if (!tx->rollback()) { a->error = "rollback failed"; return NULL; }
    return NULL;
}

// TEST A: ConcurrentTxnsStayPinned  (the most important check)
//
// kTxns overlapping transactions, each on its own bthread, on the POOLED
// channel with the pool capped SMALL (4) relative to kTxns (8) so they contend
// AND force new-connection creation.  Each txn runs 4 statements into its OWN
// scratch table and snapshots tx->GetSocketId() before every statement.
// Asserts (on the main thread, after join):
//   * INTRA-txn: the SocketId is CONSTANT across a transaction's statements
//     (the txn is pinned to one connection);
//   * INTER-txn: the live SocketIds are DISTINCT across the concurrent txns
//     (no two simultaneously-open transactions share a pooled connection).
// The whole thing loops a few times to shake scheduling.
TEST_F(MysqlPoolConcurrencyTest, ConcurrentTxnsStayPinned) {
    const int kTxns = 8;
    SetPoolCap(4);  // SMALL vs kTxns=8: contend + force new connections.

    std::string err;
    brpc::MysqlResponse resp;
    std::vector<std::string> tables(kTxns);
    for (int w = 0; w < kTxns; ++w) {
        tables[w] = butil::string_printf("pool_conc_pin_w%d", w);
        ASSERT_TRUE(RunPlain(_channel, "DROP TABLE IF EXISTS " + tables[w],
                             &resp, &err)) << err;
        ASSERT_TRUE(RunPlain(_channel,
                             "CREATE TABLE " + tables[w] +
                                 " (id INT PRIMARY KEY, name VARCHAR(32)) "
                                 "ENGINE=InnoDB",
                             &resp, &err)) << err;
    }

    for (int loop = 0; loop < 10; ++loop) {
        std::vector<PinnedTxnWorkerArgs> args(kTxns);
        for (int w = 0; w < kTxns; ++w) {
            args[w].channel = &_channel;
            args[w].table = tables[w];
            // Unique per worker AND per loop so rows never collide.
            args[w].id = 120000 + loop * 100 + w;
        }

        std::vector<bthread_t> threads(kTxns);
        for (int w = 0; w < kTxns; ++w) {
            ASSERT_EQ(0, bthread_start_background(&threads[w], NULL,
                                                  PinnedTxnWorker, &args[w]));
        }
        for (int w = 0; w < kTxns; ++w) {
            bthread_join(threads[w], NULL);
        }

        // No worker errored.
        for (int w = 0; w < kTxns; ++w) {
            ASSERT_TRUE(args[w].error.empty())
                << "loop " << loop << " txn " << w << ": " << args[w].error;
            ASSERT_EQ(4u, args[w].socket_ids.size()) << "loop " << loop;
        }

        // INTRA-txn pinning: one constant non-zero SocketId per transaction.
        std::vector<brpc::SocketId> live(kTxns);
        for (int w = 0; w < kTxns; ++w) {
            const brpc::SocketId sid = args[w].socket_ids[0];
            EXPECT_NE(0u, sid) << "loop " << loop << " txn " << w;
            for (size_t s = 1; s < args[w].socket_ids.size(); ++s) {
                EXPECT_EQ(sid, args[w].socket_ids[s])
                    << "loop " << loop << " txn " << w << " stmt " << s
                    << ": transaction was NOT pinned to one connection";
            }
            live[w] = sid;
        }

        // INTER-txn isolation: all live (simultaneously-open) SocketIds distinct.
        for (int i = 0; i < kTxns; ++i) {
            for (int j = i + 1; j < kTxns; ++j) {
                EXPECT_NE(live[i], live[j])
                    << "loop " << loop << ": txns " << i << " and " << j
                    << " shared a pooled connection (sid=" << live[i] << ")";
            }
        }
    }

    for (int w = 0; w < kTxns; ++w) {
        RunPlain(_channel, "DROP TABLE IF EXISTS " + tables[w], &resp, &err);
    }
}

// TEST B workers: an aborting-transaction mix.  Mode 0 explicitly rollback()s
// after an INSERT; mode 1 simply drops the MysqlTransactionUniquePtr WITHOUT
// commit, so its destructor auto-rollbacks (see MysqlTransaction::~ in
// mysql_transaction.h).  Either way the insert must NOT survive.
struct AbortWorkerArgs {
    brpc::Channel* channel;
    std::string table;   // shared abort table
    int id;              // unique id this worker inserts then aborts
    int mode;            // 0 == explicit rollback, 1 == drop -> dtor rollback
    std::string error;
};

void* AbortWorker(void* p) {
    AbortWorkerArgs* a = static_cast<AbortWorkerArgs*>(p);
    a->error.clear();
    {
        brpc::MysqlTransactionUniquePtr tx = brpc::NewMysqlTransaction(
            *a->channel, brpc::MysqlTransactionOptions());
        if (tx == NULL) { a->error = "NewMysqlTransaction NULL"; return NULL; }

        brpc::MysqlResponse resp;
        if (!RunInTx(*a->channel, tx.get(),
                     butil::string_printf("INSERT INTO %s VALUES (%d, 'quill')",
                                          a->table.c_str(), a->id),
                     &resp, &a->error)) return NULL;
        if (resp.reply(0).is_error()) {
            a->error = "INSERT err: " + resp.reply(0).error().msg().as_string();
            return NULL;
        }
        if (a->mode == 0) {
            if (!tx->rollback()) { a->error = "explicit rollback failed"; return NULL; }
        }
        // mode 1: fall off the end of this scope -> tx dtor auto-rollbacks.
    }
    return NULL;
}

// TEST B: ConcurrentTxnAbortAndAutoRollback
//
// Under the same small-pool contended setup, run a concurrent MIX of
// explicitly-rolled-back transactions and dropped-without-commit transactions
// (whose dtor auto-rollbacks).  Exercises the reserve -> return / auto-rollback
// path under contention (the UAF/leak-prone path).  After join, from a fresh
// non-tx connection, assert NONE of the aborted inserts are visible, workers
// saw no errors, and the channel is still healthy (a final pooled SELECT
// succeeds -> reserved connections were returned to the pool cleanly).
TEST_F(MysqlPoolConcurrencyTest, ConcurrentTxnAbortAndAutoRollback) {
    const int kWorkersB = 8;
    SetPoolCap(4);  // SMALL vs 8 workers: contend + force new connections.

    const std::string t = "pool_conc_abort";
    std::string err;
    brpc::MysqlResponse resp;
    ASSERT_TRUE(RunPlain(_channel, "DROP TABLE IF EXISTS " + t, &resp, &err)) << err;
    ASSERT_TRUE(RunPlain(_channel,
                         "CREATE TABLE " + t +
                             " (id INT PRIMARY KEY, name VARCHAR(32)) "
                             "ENGINE=InnoDB",
                         &resp, &err)) << err;

    for (int loop = 0; loop < 6; ++loop) {
        std::vector<AbortWorkerArgs> args(kWorkersB);
        for (int w = 0; w < kWorkersB; ++w) {
            args[w].channel = &_channel;
            args[w].table = t;
            args[w].id = 130000 + loop * 100 + w;  // unique per loop+worker
            args[w].mode = w % 2;  // half explicit rollback, half dtor rollback
        }

        std::vector<bthread_t> threads(kWorkersB);
        for (int w = 0; w < kWorkersB; ++w) {
            ASSERT_EQ(0, bthread_start_background(&threads[w], NULL,
                                                  AbortWorker, &args[w]));
        }
        for (int w = 0; w < kWorkersB; ++w) {
            bthread_join(threads[w], NULL);
        }

        for (int w = 0; w < kWorkersB; ++w) {
            EXPECT_TRUE(args[w].error.empty())
                << "loop " << loop << " worker " << w << " (mode "
                << args[w].mode << "): " << args[w].error;
        }

        // Effects discarded: not a single aborted insert is visible from a
        // fresh (non-tx) pooled connection.
        ASSERT_TRUE(RunPlain(_channel,
                             "SELECT COUNT(*) FROM " + t, &resp, &err)) << err;
        long long visible = -1;
        ASSERT_EQ(1, ResultRowCount(resp)) << "loop " << loop;
        ASSERT_TRUE(FieldToLongLong(resp.reply(0).next().field(0), &visible));
        EXPECT_EQ(0, visible)
            << "loop " << loop << ": " << visible
            << " aborted/dropped insert(s) leaked into the table";
    }

    // Channel still healthy: a final simple SELECT on a fresh pooled request
    // succeeds -> the reserved connections were returned to the pool cleanly.
    ASSERT_TRUE(RunPlain(_channel, "SELECT 1", &resp, &err)) << err;
    EXPECT_EQ(1, ResultRowCount(resp));

    RunPlain(_channel, "DROP TABLE IF EXISTS " + t, &resp, &err);
}

// TEST C support: a transaction that RESERVES a pooled connection (pulls it out
// of the pool) and holds it for the lifetime of the worker, so that the
// prepared statement S is forced onto connections that may not have its
// server-side stmt_id cached.
struct ReserveWorkerArgs {
    brpc::Channel* channel;
    std::string error;
};

void* ReserveWorker(void* p) {
    ReserveWorkerArgs* a = static_cast<ReserveWorkerArgs*>(p);
    a->error.clear();
    // Open a few short transactions in series; each pulls a pooled connection
    // out (reserve) and returns it (rollback), churning which sockets are in
    // the pool while S is being executed concurrently.
    for (int k = 0; k < 8; ++k) {
        brpc::MysqlTransactionUniquePtr tx = brpc::NewMysqlTransaction(
            *a->channel, brpc::MysqlTransactionOptions());
        if (tx == NULL) { a->error = "reserve: NewMysqlTransaction NULL"; return NULL; }
        brpc::MysqlResponse resp;
        if (!RunInTx(*a->channel, tx.get(), "SELECT 1", &resp, &a->error)) return NULL;
        if (!tx->rollback()) { a->error = "reserve: rollback failed"; return NULL; }
    }
    return NULL;
}

// Execute a shared prepared statement S with a fresh INT param and verify the
// bound value comes back correctly.  S is shared across bthreads, so it lands
// on whatever pooled connection is currently free.
struct StmtExecWorkerArgs {
    brpc::Channel* channel;
    brpc::MysqlStatement* stmt;  // shared prepared statement S
    int base;                    // param seed (unique per worker)
    std::string error;
};

void* StmtExecWorker(void* p) {
    StmtExecWorkerArgs* a = static_cast<StmtExecWorkerArgs*>(p);
    a->error.clear();
    for (int k = 0; k < 12; ++k) {
        const int32_t v = a->base + k;
        brpc::MysqlRequest req(a->stmt);
        if (!req.AddParam(v)) { a->error = "AddParam failed"; return NULL; }
        brpc::MysqlResponse resp;
        brpc::Controller cntl;
        a->channel->CallMethod(NULL, &cntl, &req, &resp, NULL);
        if (cntl.Failed()) { a->error = "rpc: " + cntl.ErrorText(); return NULL; }
        if (resp.reply_size() < 1 || !resp.reply(0).is_resultset() ||
            resp.reply(0).row_count() != 1u) {
            a->error = "bad resultset for S"; return NULL;
        }
        long long got = 0;
        if (!FieldToLongLong(resp.reply(0).next().field(0), &got) || got != v) {
            a->error = butil::string_printf(
                "S returned wrong value (got %lld want %d)", got, v);
            return NULL;
        }
    }
    return NULL;
}

// TEST C: PreparedRePreparesWhenConnectionStolen
//
// MECHANISM: a server-side prepared statement id is per-(SocketId, fd_version)
// -- it is meaningful only on the exact connection that ran COM_STMT_PREPARE
// (see MysqlStatement::StatementId(SocketId)/SetStatementId(SocketId,...) in
// mysql_statement.h, and the fd_version/SocketId ABA discussion in this file's
// header).  When a shared MysqlStatement S lands on a pooled connection that
// does NOT have S's stmt_id cached, brpc must transparently issue a fresh
// COM_STMT_PREPARE on that connection before COM_STMT_EXECUTE, and the bound
// result must still be correct.
//
// SETUP: pool capped at 2.  S = "SELECT CAST(? AS SIGNED) AS v" is created and
// executed once (caching its stmt_id on whatever connection it first landed
// on).  Then background bthreads open transactions that RESERVE the two pooled
// connections (pulling S's original connection out), while other bthreads keep
// executing S with fresh params.  Looped with churn so S repeatedly hits
// connections it was not prepared on.  Every execute must still return the
// correct bound value; per-worker errors are recorded and asserted empty after
// join.
TEST_F(MysqlPoolConcurrencyTest, PreparedRePreparesWhenConnectionStolen) {
    SetPoolCap(2);  // tiny pool: 2 sockets, easy to "steal" via reserving txns.

    std::string err;
    brpc::MysqlResponse resp;

    brpc::MysqlStatementUniquePtr S =
        brpc::NewMysqlStatement(_channel, "SELECT CAST(? AS SIGNED) AS v");
    ASSERT_TRUE(S != NULL);
    ASSERT_EQ(1u, S->param_count());

    // Execute S once to cache its stmt_id on whatever connection it lands on.
    {
        brpc::MysqlRequest req(S.get());
        ASSERT_TRUE(req.AddParam((int32_t)140000));
        brpc::Controller cntl;
        _channel.CallMethod(NULL, &cntl, &req, &resp, NULL);
        ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
        ASSERT_TRUE(resp.reply_size() >= 1 && resp.reply(0).is_resultset());
        long long got = 0;
        ASSERT_TRUE(FieldToLongLong(resp.reply(0).next().field(0), &got));
        ASSERT_EQ(140000, got);
    }

    const int kExecutors = 4;
    const int kReservers = 3;
    for (int loop = 0; loop < 30; ++loop) {
        std::vector<StmtExecWorkerArgs> exec_args(kExecutors);
        std::vector<ReserveWorkerArgs> res_args(kReservers);
        std::vector<bthread_t> exec_threads(kExecutors);
        std::vector<bthread_t> res_threads(kReservers);

        for (int w = 0; w < kReservers; ++w) {
            res_args[w].channel = &_channel;
            ASSERT_EQ(0, bthread_start_background(&res_threads[w], NULL,
                                                  ReserveWorker, &res_args[w]));
        }
        for (int w = 0; w < kExecutors; ++w) {
            exec_args[w].channel = &_channel;
            exec_args[w].stmt = S.get();
            // Unique param ranges per worker+loop so a wrong value is unambiguous.
            exec_args[w].base = 141000 + loop * 1000 + w * 100;
            ASSERT_EQ(0, bthread_start_background(&exec_threads[w], NULL,
                                                  StmtExecWorker, &exec_args[w]));
        }
        for (int w = 0; w < kReservers; ++w) {
            bthread_join(res_threads[w], NULL);
        }
        for (int w = 0; w < kExecutors; ++w) {
            bthread_join(exec_threads[w], NULL);
        }

        for (int w = 0; w < kReservers; ++w) {
            EXPECT_TRUE(res_args[w].error.empty())
                << "loop " << loop << " reserver " << w << ": " << res_args[w].error;
        }
        for (int w = 0; w < kExecutors; ++w) {
            EXPECT_TRUE(exec_args[w].error.empty())
                << "loop " << loop << " executor " << w << ": " << exec_args[w].error;
        }
    }
}

}  // namespace