增强你的 sysbench - 给 TiDB 添加自定义测试

对于从事数据库相关的同学来说，对数据库进行性能测试是一个永远绕不开的话题。这个世界上有很多的数据库性能测试工具，而 sysbench 可以算是大家用的最多的之一。

根据官网的介绍，sysbench 是一个基于 Luajit 的，多线程的脚本化的性能测试工具，它不光能测试数据库的性能，也可以测试文件 I/O，CPU 这些，不过这里，我们重点来聊聊数据库的测试。在 sysbench 里面，对于数据库的测试，用的最多的就是 OLTP 相关的脚本，譬如 oltp_point_select，oltp_update_index。通常，OLTP 的脚本在绝大多数时候都够用了，但程序员就是这么的不满足，尤其是对我们这些开发数据库的程序员来说，我们不光希望能测试 OLTP 相关的 workload，还希望能通过 sysbench 测试更多的 workload。

幸运的是，随着 sysbench 1.0 的发布，它开始支持自定义的脚本来进行 benchmark，那么我们要做的事情就很简单了，只需要干两件事情：

1. 学习 Lua 这门脚本语言。
2. 给 sysbench 写自己的插件。

对于 Lua 语言，这里不做过多讨论，即使你不会，也不用特别担心，它是一门非常容易学习的语言。这里我们来谈谈如何给 sysbench 写插件。

基本框架

这里，我们来实现一个 bank transfer 的 benchmark，首先我们创建一个 bank_transfer.lua 文件，在文件头写上

#!/usr/bin/env sysbench

然后基本的框架：

function thread_init() 
    print(string.format("start thread %d", sysbench.tid))
end

function thread_done()
    print(string.format("stop thread %d", sysbench.tid))
end

function event() 
end

在上面的文件里面，我们实现了三个函数，thread_init 和 thread_done 非常的直观，就是 sysbench 在测试线程启动和结束的时候调用，而 event 则是我们实际执行测试的地方，这里是空的。然后我们运行这个测试程序，你可以暂时忽略命令行里面的 MySQL 参数，这个后面我们会实际连接到 MySQL 进行测试。

sysbench --report-interval=1 --time=20 --threads=4 --mysql-host=127.0.0.1 --mysql-port=4000 --mysql-user=root --mysql-db=sbtest --db-driver=mysql bank_transfer run

start thread 0
start thread 2
start thread 1
start thread 3
Threads started!

stop thread 0
stop thread 3
stop thread 1
stop thread 2

传递参数

上面例子里面的参数其实是 sysbench 自己需要的参数，对于我们的测试程序来说，它也可能需要一些参数，所以我们在脚本里面定义自己的参数，如下我们定义了两个参数，table-size 和 tables，用来告诉测试脚本，我们希望创建多少张表以及每张表里面有多少数据。

sysbench.cmdline.options = {
    table_size =
        {"Number of rows per table", 10000},
    tables =
        {"Number of tables", 1}
}

function thread_init() 
    print(sysbench.opt.table_size)
end

启动的时候我们就可以指定自己的参数了，如下：

sysbench bank_transfer run --tables=16 --table-size=1000000

连接数据库

因为我们是要对数据库进行测试，所以首先我们需要跟数据库建立起连接，在脚本里面写上：

function thread_init() 
    drv = sysbench.sql.driver()
    con = drv:connect()
end

function thread_done()
    con:disconnect()
end

function event() 
    con:query("select 1")
end

上面我们在线程初始化的时候跟数据库建立了连接，在结束的时候关闭了连接，而在 event 里面，则是执行了 select 1 这个操作。重新执行脚本，我们会看到如下输出：

Threads started!

[ 1s ] thds: 4 tps: 24940.85 qps: 24940.85 (r/w/o: 24940.85/0.00/0.00) lat (ms,95%): 0.27 err/s: 0.00 reconn/s: 0.00
[ 2s ] thds: 4 tps: 28102.97 qps: 28102.97 (r/w/o: 28102.97/0.00/0.00) lat (ms,95%): 0.24 err/s: 0.00 reconn/s: 0.00

Prepare 和 cleanup

在开始测试之前，我们首先要导入数据，在 sysbench 里面，这个是通过 prepare 来完成的，首先我们定义好 prepare 函数：

function cmd_prepare()
    local drv = sysbench.sql.driver()
    local con = drv:connect()

    for i = sysbench.tid % sysbench.opt.threads + 1, sysbench.opt.tables, sysbench.opt
        .threads do 
        create_table(drv, con, i) 
    end
end

sysbench.cmdline.commands = {
    prepare = {cmd_prepare, sysbench.cmdline.PARALLEL_COMMAND}
}

上面我们定义了 cmd_prepare，并且告诉 sysbench 这个是 PARALLEL_COMMAND 类型，也就是 sysbench 会并发的调用 prepare。而在 cmd_prepare 里面，我们也是让每个线程负责给不同的 table 导入数据。

我们在 create_table 里面创建表结构，以及使用 bulk_insert_* 相关的接口来导入数据：

function create_table(drv, con, table_num)
    print(string.format("Creating table 'account%d'...", table_num))

    local query = string.format([[
CREATE TABLE account%d(
  id INTEGER NOT NULL,
  balance INTEGER DEFAULT '1000' NOT NULL,
  PRIMARY KEY (id)
)]], table_num)

    con:query(query)

    if (sysbench.opt.table_size > 0) then
        print(string.format("Inserting %d records into 'account%d'",
                            sysbench.opt.table_size, table_num))
    end

    query = "INSERT INTO account" .. table_num .. "(id, balance) VALUES"

    con:bulk_insert_init(query)

    for i = 1, sysbench.opt.table_size do
        query = string.format("(%d, %d)", i, 1000)

        con:bulk_insert_next(query)
    end

    con:bulk_insert_done()
end

我们能 prepare，自然也会有对应的 cleanup:

function cleanup()
    local drv = sysbench.sql.driver()
    local con = drv:connect()

    for i = 1, sysbench.opt.tables do
        print(string.format("Dropping table 'account%d'...", i))
        con:query("DROP TABLE IF EXISTS account" .. i)
    end
end

当定义好上面这些函数之后，我们就可以使用 sysbench bank_transfer prepare 以及 sysbench bank_transfer cleanup 来导入或者清理数据了。

Transfer

好了，现在到了最激动人心的时刻，我们开始写真正的测试逻辑。我们这里要模拟的是 transfer，代码如下：

function event() 
    local from = get_id()
    local to = get_id()
    local table_num = get_table_num()
    local amount = sysbench.rand.default(1, 100)
    while(from == to)
    do
        to = get_id()
    end

    con:query("BEGIN")

    local rs = con:query(string.format([[
SELECT id, balance FROM account%d WHERE id IN (%d, %d) FOR UPDATE
]], table_num, from, to))

    assert(rs.nrows == 2)

    local row_from = rs:fetch_row()
    local row_to = rs:fetch_row()

    if row_from[1] ~= from then
        row_from, row_to = row_to, row_from
    end 

    if row_from[2] - amount < 0 then 
        con:query("ROLLBACK")
        return 
    end

    con:query(string.format([[
UPDATE account%d SET balance = balance - %d WHERE id = %d
]], table_num, amount, from))

    con:query(string.format([[
UPDATE account%d SET balance = balance + %d WHERE id = %d
]], table_num, amount, to))

    con:query("COMMIT") 
end

可以看到，逻辑还是非常简单的，主要流程是：

1. 开启事务
2. 随机选择两个账户 from 和 to
3. 查询两个账户的余额
4. 如果 from 的余额不够，转账失败，回滚事务
5. 执行转账操作，from 减去 amount，to 增加 amount
6. 提交事务

运行测试，可以看到如下输出：

[ 1s ] thds: 4 tps: 661.92 qps: 3325.56 (r/w/o: 665.91/1331.82/1327.83) lat (ms,95%): 10.09 err/s: 0.00 reconn/s: 0.00
[ 2s ] thds: 4 tps: 678.85 qps: 3394.23 (r/w/o: 678.85/1357.69/1357.69) lat (ms,95%): 7.98 err/s: 0.00 reconn/s: 0.00
[ 3s ] thds: 4 tps: 894.98 qps: 4465.85 (r/w/o: 891.96/1783.93/1789.95) lat (ms,95%): 5.18 err/s: 0.00 reconn/s: 0.00

总结

上面通过一个简单的例子，告诉大家如何在 sysbench 里面写自己的测试，我个人认为，作为一个非常通用的测试框架，会有越来越多的开发者给 sysbench 添加新的测试用例，譬如 Percona 已经添加了 tpcc，blob，我自己后续也会尝试在 sysbench 里面加入更多的测试 case，来对 TiDB 进行各种的性能测试。