wehub-resource-sync
145 lines
3.5 KiB
Lua
145 lines
3.5 KiB
Lua
local stress_generator = {}
|
|
stress_generator.__index = stress_generator
|
|
|
|
|
|
local cjson = require "cjson"
|
|
local pl_file = require "pl.file"
|
|
local pl_path = require "pl.path"
|
|
local uuid = require "resty.jit-uuid"
|
|
|
|
|
|
local fmt = string.format
|
|
local tmp_root = os.getenv("TMPDIR") or "/tmp"
|
|
|
|
|
|
-- we need this to get random UUIDs
|
|
math.randomseed(os.time())
|
|
|
|
|
|
local attack_cmds = {
|
|
["http"] = "GET http://%s:%d%s",
|
|
}
|
|
|
|
|
|
function stress_generator.is_running(self)
|
|
if self.finish_time == nil or self.finish_time <= os.time() then
|
|
return false
|
|
end
|
|
|
|
return true
|
|
end
|
|
|
|
|
|
function stress_generator.get_results(self)
|
|
if self.results ~= nil then
|
|
return self.results
|
|
end
|
|
|
|
if self.results_filename == nil then
|
|
return nil, "stress_generator was not run yet"
|
|
end
|
|
|
|
if stress_generator:is_running() then
|
|
return nil, "stress_generator results not available yet"
|
|
end
|
|
|
|
local report_cmd = fmt("vegeta report -type=json %s 2>&1", self.results_filename)
|
|
local report_pipe = io.popen(report_cmd)
|
|
local output = report_pipe:read('*all')
|
|
report_pipe:close()
|
|
|
|
if pl_path.exists(self.results_filename) then
|
|
pl_file.delete(self.results_filename)
|
|
end
|
|
|
|
|
|
local vegeta_results = cjson.decode(output)
|
|
local results = {
|
|
["successes"] = 0,
|
|
["remote_failures"] = 0,
|
|
["proxy_failures"] = 0,
|
|
["failures"] = 0,
|
|
}
|
|
|
|
vegeta_results.status_codes = vegeta_results.status_codes or {}
|
|
|
|
for status, count in pairs(vegeta_results.status_codes) do
|
|
if status == "200" then
|
|
results.successes = count
|
|
elseif status == "502" or status == "504" then
|
|
results.remote_failures = results.remote_failures + count
|
|
results.failures = results.failures + count
|
|
elseif status == "500" or status == "503" then
|
|
results.proxy_failures = results.proxy_failures + count
|
|
results.failures = results.failures + count
|
|
else
|
|
results.failures = results.failures + count
|
|
end
|
|
end
|
|
|
|
self.results = results
|
|
|
|
if self.debug then
|
|
-- show pretty results
|
|
local report_cmd = fmt("vegeta report %s 2>&1", self.results_filename)
|
|
local report_pipe = io.popen(report_cmd)
|
|
local output = report_pipe:read('*all')
|
|
report_pipe:close()
|
|
print(output)
|
|
end
|
|
|
|
return self.results
|
|
end
|
|
|
|
|
|
function stress_generator.run(self, uri, headers, duration, rate)
|
|
if stress_generator:is_running() then
|
|
return nil, "already running"
|
|
end
|
|
|
|
self.results_filename = fmt("%s/vegeta_%s", tmp_root, uuid())
|
|
|
|
duration = duration or 1
|
|
rate = rate or 100
|
|
local attack_cmd = fmt(attack_cmds[self.protocol], self.host, self.port, uri)
|
|
local req_headers = ""
|
|
|
|
for name, value in pairs(headers) do
|
|
req_headers = fmt("-header=%s:%s %s", name, value, req_headers)
|
|
end
|
|
|
|
local vegeta_cmd = fmt(
|
|
"echo %s | vegeta attack %s -rate=%d -duration=%ds -workers=%d -timeout=5s -output=%s",
|
|
attack_cmd, req_headers, rate, duration, self.workers, self.results_filename)
|
|
|
|
self.pipe = io.popen(vegeta_cmd)
|
|
-- we will rely on vegeta's duration
|
|
self.finish_time = os.time() + duration
|
|
end
|
|
|
|
|
|
function stress_generator.new(protocol, host, port, workers, debug)
|
|
if io.popen == nil then
|
|
error("stress_generator is not supported in this platform", 2)
|
|
end
|
|
|
|
local self = setmetatable({}, stress_generator)
|
|
|
|
protocol = protocol or "http"
|
|
|
|
if protocol ~= "http" then
|
|
error("stress_generator supports only http")
|
|
end
|
|
|
|
self.debug = debug == true
|
|
self.host = host or "127.0.0.1"
|
|
self.port = port or "80"
|
|
self.protocol = protocol
|
|
self.workers = workers or 10
|
|
|
|
return self
|
|
end
|
|
|
|
|
|
return stress_generator
|