wrk - lua 进阶

wrk是一个小型高性能的接口压力测试的小工具，最近学习了一下，对于开发来说还是比较好用的，易上手，可编程扩展，使用lua脚本可以对其进行一下自定义，所以这里就对wrk中使用lua进行探究

在wrk中是通过自定义相应的lua方法达到改变wrk行为的目的，wrk的执行分为三个阶段：启动阶段（setup）、运行阶段(running)、结束阶段(done)，每个测试线程，都拥有独立的lua 运行环境。

wrk执行流程

1. 启动阶段:

function setup(thread)

在脚本文件中实现 setup 方法，wrk 就会在测试线程已经初始化，但还没有启动的时候调用该方法。wrk会为每一个测试线程调用一次 setup 方法，并传入代表测试线程的对象thread 作为参数。setup 方法中可操作该thread 对象，获取信息、存储信息、甚至关闭该线程。

thread.addr             -- get or set the thread's server address
thread:get(name)        -- get the value of a global in the thread's env
thread:set(name, value) -- set the value of a global in the thread's env
thread:stop()           -- stop the thread

运行阶段：

function init(args)
function delay()
function request()
function response(status, headers, body)

init(args):
由测试线程调用，只会在进入运行阶段时，调用一次。支持从启动 wrk 的命令中，获取命令行参数
delay()：
每次发送请求之前调用，可以在这里定制延迟时间，通过返回值（单位毫秒）如:return 1000，即延迟一秒
request():
每次发送请求之前调用，可以对每一次的请求做一些自定义的操作，但是不要在该方法中做耗时的操作
response(status, headers, body):
在每次收到一个响应时被调用，为提升性能，如果没有定义该方法，为了提升效率，那么wrk不会解析 headers 和 body

结束阶段：

function done(summary, latency, requests)

done()
该方法和setup方法一样，只会被调用一次，整个测试完后执行，在定的参数中获取压测结果，生成定制化的测试报告

下面是官方对Lua API的说明

The public Lua API consists of a global table and a number of global functions:
  wrk = {
    scheme  = "http",
    host    = "localhost",
    port    = nil,
    method  = "GET",
    path    = "/",
    headers = {},
    body    = nil,
    thread  = <userdata>,
  }

  function wrk.format(method, path, headers, body)

    wrk.format returns a HTTP request string containing the passed parameters
    merged with values from the wrk table.

  function wrk.lookup(host, service)

    wrk.lookup returns a table containing all known addresses for the host
    and service pair. This corresponds to the POSIX getaddrinfo() function.

  function wrk.connect(addr)

    wrk.connect returns true if the address can be connected to, otherwise
    it returns false. The address must be one returned from wrk.lookup().

  The following globals are optional, and if defined must be functions:

    global setup    -- called during thread setup
    global init     -- called when the thread is starting
    global delay    -- called to get the request delay
    global request  -- called to generate the HTTP request
    global response -- called with HTTP response data
    global done     -- called with results of run

Setup

  function setup(thread)

  The setup phase begins after the target IP address has been resolved and all
  threads have been initialized but not yet started.

  setup() is called once for each thread and receives a userdata object
  representing the thread.

    thread.addr             - get or set the thread's server address
    thread:get(name)        - get the value of a global in the thread's env
    thread:set(name, value) - set the value of a global in the thread's env
    thread:stop()           - stop the thread

  Only boolean, nil, number, and string values or tables of the same may be
  transfered via get()/set() and thread:stop() can only be called while the
  thread is running.

Running

  function init(args)
  function delay()
  function request()
  function response(status, headers, body)

  The running phase begins with a single call to init(), followed by
  a call to request() and response() for each request cycle.

  The init() function receives any extra command line arguments for the
  script which must be separated from wrk arguments with "--".

  delay() returns the number of milliseconds to delay sending the next
  request.

  request() returns a string containing the HTTP request. Building a new
  request each time is expensive, when testing a high performance server
  one solution is to pre-generate all requests in init() and do a quick
  lookup in request().

  response() is called with the HTTP response status, headers, and body.
  Parsing the headers and body is expensive, so if the response global is
  nil after the call to init() wrk will ignore the headers and body.

Done

  function done(summary, latency, requests)

  The done() function receives a table containing result data, and two
  statistics objects representing the per-request latency and per-thread
  request rate. Duration and latency are microsecond values and rate is
  measured in requests per second.

  latency.min              -- minimum value seen
  latency.max              -- maximum value seen
  latency.mean             -- average value seen
  latency.stdev            -- standard deviation
  latency:percentile(99.0) -- 99th percentile value
  latency(i)               -- raw value and count

  summary = {
    duration = N,  -- run duration in microseconds
    requests = N,  -- total completed requests
    bytes    = N,  -- total bytes received
    errors   = {
      connect = N, -- total socket connection errors
      read    = N, -- total socket read errors
      write   = N, -- total socket write errors
      status  = N, -- total HTTP status codes > 399
      timeout = N  -- total request timeouts
    }
  }

三个阶段综合小例子

例一

-- 启动阶段 （每个线程执行一次）
function setup(thread)
   print("----------启动阶段----------------")
   print("setup",thread)
   print("setup",thread.addr)
end


-- 运行阶段 （该方法init每个线程执行一次）
function init(args)
   print("-----------运行阶段---------------")
   print("init",args)
end


-- 这个三个方法每个请求都会调用一次
function delay()
 print("delay")
 -- 设置延迟990ms
 return 990
end

function request()
  print("request")
  -- 这个方法必须要有返回，不然会出错
  return wrk.request()
end

function response(status, headers, body)
  print("response",status,headers)
end



-- 结束阶段 
function done(summary, latency, requests)
   print("-----------结束阶段---------------")
  print("done",summary,latency,requests)
end

运行结果 运行结果

从运行结果中可以看出，先是调用setup启动，然后进行init初始化，然后再调用request获取请求内容，在调用delay获取延迟时间，这里因为delay比较多，所以当response被调用后才开始第二轮

例二

local threads = {}


-- 启动阶段 （每个线程执行一次）
function setup(thread)
    print("----------启动阶段----------------")
    table.insert(threads,thread)
end


-- 运行阶段 （该方法init每个线程执行一次）
function init(args)
    print("-----------运行阶段---------------")
    print("init:threads中的元素个数:",#threads)
end


-- 这个三个方法每个请求都会调用一次
function delay()
  print("delay:threads中的元素个数:",#threads)
  -- 设置延迟990ms
  return 990
end      
          
function request()
    -- 这个方法必须要有返回，不然会出错
    print("request:threads中的元素个数:",#threads)
    return wrk.request()
end   

function response(status, headers, body)  
    print("response:threads中的元素个数:",#threads)
end 

    
    
-- 结束阶段 
function done(summary, latency, requests)
    print("-----------结束阶段---------------")
    print("done:threads中的元素个数:",#threads)
end

例二运行结果

当一个局部变量table内部元素在启动阶段（setup）发生改变时，在运行阶段调用的方法是直接无法获取的，在结束阶段（done）是可以的

例三

    print("----------启动阶段----------------")
    table.insert(threads,thread)
end


-- 运行阶段 （该方法init每个线程执行一次）
function init(args)
    print("-----------运行阶段---------------")
    print("init:threads中的元素个数:",#threads)
    counter = 222
end


-- 这个三个方法每个请求都会调用一次
function delay()
  -- 设置延迟990ms
  print("delay:threads中的元素个数:",#threads)
  print("delay:counter:",counter)
  return 990
end

function request()
    -- 这个方法必须要有返回，不然会出错
    print("request:threads中的元素个数:",#threads)
    print("requset:counter:",counter) 
    return wrk.request() 
end
             
function response(status, headers, body)
    print("response:threads中的元素个数:",#threads)
    print("response:counter:",counter) 
end 
    

    
-- 结束阶段 
function done(summary, latency, requests)
    print("-----------结束阶段---------------")
    print("done:threads中的元素个数:",#threads)
    print("done:counter:",counter)
end

例三运行结果

同样在运行阶段对局部变量的改变也在结束阶段获取不到

例四

local counter = 1
local threads = {}


-- 启动阶段 （每个线程执行一次）
function setup(thread)
    print("----------启动阶段----------------")
    table.insert(threads,thread)
    print("setup:thrads地址:",threads)
    counter = 111
end


-- 运行阶段 （该方法init每个线程执行一次）
function init(args)
    print("-----------运行阶段---------------")
    print("init:threads中的元素个数:",#threads,threads)
    counter = 222
end


-- 这个三个方法每个请求都会调用一次
function delay()
  -- 设置延迟990ms
  print("delay:threads中的元素个数:",#threads,threads)
  print("delay:counter:",counter)
  return 990
end     

function request()
    -- 这个方法必须要有返回，不然会出错
    print("request:threads中的元素个数:",#threads,threads)
    print("requset:counter:",counter)
    return wrk.request()
end 

function response(status, headers, body)   
    print("response:threads中的元素个数:",#threads,threads)
    print("response:counter:",counter)
end 



-- 结束阶段 
function done(summary, latency, requests)
    print("-----------结束阶段---------------")
    print("done:threads中的元素个数:",#threads,threads)
    print("done:counter:",counter)
end

例四运行结果

从运行结果发现其实改变无法察觉是因为他们并不是同一个变量，结束阶段打印的threads地址和运行阶段打印的地址是不一样的，但是启动阶段和结束阶段是一样的

将例四中的threads变量改为全局变量后运行

将例四中的threads变量改为全局变量后运行

运行结果是一样的，这里我得出一个结论，前面说每个测试线程，都拥有独立的lua 运行环境，这个独立的环境是运行阶段用不同的线程体现，首先是有一个主线程获取命令行中的参数信息，然后解析-t 1后得知要创建一个线程thread，让后会创建一个线程在主线程中将创建的线程thread传入到以主线程环境的setup，当时间-d过了，主线程被唤醒，主线程关闭创建的线程，然后在主线程环境中执行 done方法，因为setup和done都是在主线环境中执行，线程上下文一样所以共享全局和外部定义的局部变量。

那么问题来了，如何传递这些变量呢？

这就轮到上面说过的thread:set和thread:get上场了，这两个方法分别是在主线程中将值设置给指定线程中的全局变量池中，用法如下例：

-- 启动阶段 （每个线程执行一次）
function setup(thread)
   local counter = 1
   table.insert(threads,thread)
   thread:set("counter",counter)
end


-- 运行阶段 （该方法init每个线程执行一次）
function init(args)
    print("-----------运行阶段---------------")
    counter = 111
end


-- 这个三个方法每个请求都会调用一次
function delay()
  -- 设置延迟990ms
  print("delay:counter:",counter)
  return 990 
end       
         
function request()
    -- 这个方法必须要有返回，不然会出错
    print("requset:counter:",counter)
    return wrk.request()
end          
      
function response(status, headers, body)
    print("response:counter:",counter) 
end 
    

    
-- 结束阶段 
function done(summary, latency, requests) 
    print("-----------结束阶段---------------")
    local thread = threads[1]
    local counter = thread:get("counter")
    print("done:counter:",counter)
end

传递运行结果

那么问题有来了，我们如何对多个线程进行同步呢?

wrk通过多线程机制使得压力效果有很大的提升，上面的测试基本都是使用一个线程，在实际使用中可能会用到多个线程，如果要用多个线程进行密码爆破测试的话我们每一个线程中的每一个请求都发送不同的密码，这个场景要怎么进行代码的编写呢？
对于学习了解过lua的朋友知道lua有一个协同的概念，下面我就用协同进行测试吧

function getPass()
    for i=1,1000 do
        coroutine.yield(i)
    end
end


function setup(thread)
    if not co then
        co = coroutine.create(getPass)
    end 
    thread:set("co",co)
end

function delay()
  return 990
end

function request()
    local status,i = coroutine.resume(co)
    local path = "/?query="..i
    return wrk.format(nil,path)
end       
         
function response(status, headers, body)  
    print("response",body)
end

协同运行结果1

function getPass()
    for i=1,1000 do
        coroutine.yield(i)
    end
end


function setup(thread)
    if not co then
        co = coroutine.wrap(getPass)
    end
    thread:set("co",co)
end

function delay()
  return 990
end

function request()
    local status,i = co()
    local path = "/?query="..i
    return wrk.format(nil,path)
end      
          
function response(status, headers, body)  
    print("response",body)
end

协同运行结果2

通过上面两个例子可以得出thread:set是有局限性的，不支持function和thread的传递，所以协同这个方式行不通了，将这两个加入到table中我就不演示了，测试了也是不行的。

下面用几种可用的方式进行说明

下面直接贴代码了：

function init(args)
    threadCount = args[1]
    local passes = getPass()
    local len = #passes / threadCount
    local startIdx = len * (id - 1) + 1
    local endIdx
    if threadCount == id then
        endIdx = #passes 
    else  
        endIdx = startIdx + len - 1
    end   
         
    print(string.format("id为:%s,%d-->%d",id,startIdx,endIdx))
    co = coroutine.create(
       function()
          for i = startIdx,endIdx do
              coroutine.yield(passes[i])
          end
       end
    )
end 
    
function delay()
  return 990
end 

function request()
    local status,i = coroutine.resume(co)
    if not i then

      wrk.thread:stop()
      return
    end
    local path = "/?query="..i
    return wrk.format(nil,path)

end

function response(status, headers, body)
    print("response",body)
end

参考文章
https://www.cnblogs.com/quanxiaoha/p/10661650.html
https://type.so/linux/lua-script-in-wrk.html