static int update_checkpoint(lua_State *lua)
{
static const char *func_name = "update_checkpoint";
lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1));
if (!lsb) {
return luaL_error(lua, "%s() invalid upvalueindex", func_name);
}
lsb_heka_sandbox *hsb = lsb_get_parent(lsb);
int result = 0;
int n = lua_gettop(lua);
switch (n) {
case 2: // async case
luaL_checktype(lua, 2, LUA_TNUMBER);
hsb->stats.pm_failures += (unsigned long long)lua_tonumber(lua, 2);
// fall thru
case 1:
luaL_checktype(lua, 1, LUA_TLIGHTUSERDATA);
result = hsb->cb.ucp(hsb->parent, lua_touserdata(lua, 1));
break;
case 0: // batch case
result = hsb->cb.ucp(hsb->parent, NULL);
break;
default:
return luaL_error(lua, "%s() invalid number of args: %d", func_name, n);
}
if (result) {
return luaL_error(lua, "%s() failed: rejected by the callback", func_name);
}
return result;
}
static char* test_api_assertion()
{
lsb_lua_sandbox *sb = lsb_create(NULL, "lua/counter.lua", "", NULL);
lsb_err_value ret = lsb_init(sb, NULL);
mu_assert(!ret, "lsb_init() received: %s", ret);
lsb_stop_sandbox(NULL);
mu_assert(lsb_destroy(NULL) == NULL, "not null");
mu_assert(lsb_usage(NULL, 0, 0) == 0, "not 0");
mu_assert(lsb_usage(sb, LSB_UT_MAX, 0) == 0, "not 0");
mu_assert(lsb_usage(sb, 0, LSB_US_MAX) == 0, "not 0");
mu_assert(strcmp(lsb_get_error(NULL), "") == 0, "not empty");
lsb_set_error(NULL, "foo");
mu_assert(lsb_get_lua(NULL) == NULL, "not null");
mu_assert(lsb_get_lua_file(NULL) == NULL, "not null");
mu_assert(lsb_get_parent(NULL) == NULL, "not null");
mu_assert(lsb_get_logger(NULL) == NULL, "not null");
mu_assert(lsb_get_state(NULL) == LSB_UNKNOWN, "not unknown");
lsb_add_function(NULL, lsb_test_write_output, "foo");
lsb_add_function(sb, NULL, "foo");
lsb_add_function(sb, lsb_test_write_output, NULL);
mu_assert(lsb_pcall_setup(NULL, "foo") == LSB_ERR_UTIL_NULL, "not null");
mu_assert(lsb_pcall_setup(sb, NULL) == LSB_ERR_UTIL_NULL, "not null");
lsb_add_function(NULL, NULL, NULL);
lsb_pcall_teardown(NULL);
lsb_terminate(NULL, NULL);
lsb_terminate(sb, NULL);
lsb_add_function(sb, lsb_test_write_output, "write_output");
e = lsb_destroy(sb);
mu_assert(!e, "lsb_destroy() received: %s", e);
return NULL;
}
static int async_checkpoint_update(lua_State* lua)
{
lua_sandbox* lsb = lua_touserdata(lua, lua_upvalueindex(1));
if (!lsb) {
return luaL_error(lua, "async_checkpoint_update() invalid upvalueindex");
}
luaL_checktype(lua, 1, LUA_TLIGHTUSERDATA);
luaL_checktype(lua, 2, LUA_TNUMBER);
size_t sequence_id = (size_t)lua_touserdata(lua, 1);
size_t failures = lua_tonumber(lua, 2);
hs_output_plugin* p = (hs_output_plugin*)lsb_get_parent(lsb);
int i = sequence_id % p->async_len;
pthread_mutex_lock(&p->cp_lock);
if (p->async_cp[i].input.id >= p->cp.input.id
&& p->async_cp[i].input.offset > p->cp.input.offset) {
p->cp.input.id = p->async_cp[i].input.id;
p->cp.input.offset = p->async_cp[i].input.offset;
}
if (p->async_cp[i].analysis.id >= p->cp.analysis.id
&& p->async_cp[i].analysis.offset > p->cp.analysis.offset) {
p->cp.analysis.id = p->async_cp[i].analysis.id;
p->cp.analysis.offset = p->async_cp[i].analysis.offset;
}
p->sb->stats.pm_failures += failures;
pthread_mutex_unlock(&p->cp_lock);
return 0;
}
static int inject_message_analysis(lua_State *lua)
{
luaL_checktype(lua, 1, LUA_TTABLE);
lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1));
if (NULL == lsb) {
return luaL_error(lua, "%s() invalid lightuserdata", im_func_name);
}
lsb_heka_sandbox *hsb = lsb_get_parent(lsb);
lua_pushstring(lua, hsb->name);
lua_setfield(lua, 1, LSB_LOGGER);
lua_pushstring(lua, hsb->hostname);
lua_setfield(lua, 1, LSB_HOSTNAME);
if (heka_encode_message_table(lsb, 1)) {
return luaL_error(lua, "%s() failed: %s", im_func_name, lsb_get_error(lsb));
}
size_t output_len = 0;
const char *output = lsb_get_output(lsb, &output_len);
if (hsb->cb.aim(hsb->parent, output, output_len) != 0) {
return luaL_error(lua, "%s() failed: rejected by the callback",
im_func_name);
}
++hsb->stats.im_cnt;
hsb->stats.im_bytes += output_len;
return 0;
}
int heka_encode_message(lua_State *lua)
{
int n = lua_gettop(lua);
bool framed = false;
switch (n) {
case 2:
luaL_checktype(lua, 2, LUA_TBOOLEAN);
framed = lua_toboolean(lua, 2);
// fall thru
case 1:
luaL_checktype(lua, 1, LUA_TTABLE);
break;
default:
return luaL_argerror(lua, n, "incorrect number of arguments");
}
lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1));
if (!lsb) {
return luaL_error(lua, "encode_message() invalid upvalueindex");
}
lsb_heka_sandbox *hsb = lsb_get_parent(lsb);
set_missing_headers(lua, 1, hsb);
lsb->output.pos = 0;
lsb_err_value ret = heka_encode_message_table(lsb, 1);
if (ret) {
const char *err = lsb_get_error(lsb);
if (strlen(err) == 0) err = ret;
return luaL_error(lua, "encode_message() failed: %s", err);
}
size_t len = 0;
const char *output = lsb_get_output(lsb, &len);
lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT] = len;
if (framed) {
char header[14] = "\x1e\x00\x08"; // up to 10 varint bytes and a \x1f
int hlen = lsb_pb_output_varint(header + 3, len) + 1;
lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT] = len + hlen + LSB_HDR_FRAME_SIZE;
header[1] = (char)hlen;
header[hlen + 2] = '\x1f';
luaL_Buffer b;
luaL_buffinit(lua, &b);
luaL_addlstring(&b, header, hlen + LSB_HDR_FRAME_SIZE);
luaL_addlstring(&b, output, len);
luaL_pushresult(&b);
} else {
lua_pushlstring(lua, output, len);
}
if (lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT]
> lsb->usage[LSB_UT_OUTPUT][LSB_US_MAXIMUM]) {
lsb->usage[LSB_UT_OUTPUT][LSB_US_MAXIMUM] =
lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT];
}
return 1;
}
static int read_message(lua_State *lua)
{
lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1));
if (NULL == lsb) {
return luaL_error(lua, "read_message() invalid lightuserdata");
}
lsb_heka_sandbox *hsb = lsb_get_parent(lsb);
if (!hsb->msg || !hsb->msg->raw.s) {
lua_pushnil(lua);
return 1;
}
return heka_read_message(lua, hsb->msg);
}
static int read_message(lua_State* lua)
{
lua_sandbox* lsb = lua_touserdata(lua, lua_upvalueindex(1));
if (!lsb) {
return luaL_error(lua, "read_message() invalid upvalueindex");
}
hs_output_plugin* p = (hs_output_plugin*)lsb_get_parent(lsb);
if (!p->matched || !p->msg) {
lua_pushnil(lua);
return 1;
}
return hs_read_message(lua, p->msg);
}
static int read_message(lua_State* lua)
{
void* luserdata = lua_touserdata(lua, lua_upvalueindex(1));
if (NULL == luserdata) {
return luaL_error(lua, "read_message() invalid lightuserdata");
}
lua_sandbox* lsb = (lua_sandbox*)luserdata;
hs_analysis_plugin* p = (hs_analysis_plugin*)lsb_get_parent(lsb);
if (!p->at->matched || !p->at->msg) {
lua_pushnil(lua);
return 1;
}
return hs_read_message(lua, p->at->msg);
}
static int inject_message(lua_State* L)
{
static unsigned char header[14];
void* luserdata = lua_touserdata(L, lua_upvalueindex(1));
if (NULL == luserdata) {
return luaL_error(L, "inject_message() invalid lightuserdata");
}
lua_sandbox* lsb = (lua_sandbox*)luserdata;
hs_analysis_plugin* p = (hs_analysis_plugin*)lsb_get_parent(lsb);
if (lua_type(L, 1) == LUA_TTABLE) {
lua_pushstring(L, p->sb->name);
lua_setfield(L, 1, "Logger");
lua_pushstring(L, p->at->plugins->cfg->hostname);
lua_setfield(L, 1, "Hostname");
lua_pushinteger(L, p->at->plugins->cfg->pid);
lua_setfield(L, 1, "Pid");
}
if (lsb_output_protobuf(lsb, 1, 0) != 0) {
return luaL_error(L, "inject_message() could not encode protobuf - %s",
lsb_get_error(lsb));
}
size_t output_len = 0;
const char* output = lsb_get_output(lsb, &output_len);
pthread_mutex_lock(&p->at->plugins->output.lock);
int len = hs_write_varint(header + 3, output_len);
int tlen = 4 + len + output_len;
++p->sb->stats.im_cnt;
p->sb->stats.im_bytes += tlen;
header[0] = 0x1e;
header[1] = (char)(len + 1);
header[2] = 0x08;
header[3 + len] = 0x1f;
fwrite(header, 4 + len, 1, p->at->plugins->output.fh);
fwrite(output, output_len, 1, p->at->plugins->output.fh);
p->at->plugins->output.offset += tlen;
if (p->at->plugins->output.offset >= (size_t)p->at->plugins->cfg->output_size) {
++p->at->plugins->output.id;
hs_open_output_file(&p->at->plugins->output);
}
pthread_mutex_unlock(&p->at->plugins->output.lock);
return 0;
}
static int inject_message_input(lua_State *lua)
{
lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1));
if (NULL == lsb) {
return luaL_error(lua, "%s() invalid lightuserdata", im_func_name);
}
const char *scp = NULL;
double ncp = NAN;
int t = lua_type(lua, 2);
switch (t) {
case LUA_TNUMBER:
ncp = lua_tonumber(lua, 2);
break;
case LUA_TSTRING:
scp = lua_tostring(lua, 2);
break;
case LUA_TNONE:
case LUA_TNIL:
break;
default:
return luaL_error(lua, "%s() unsupported checkpoint type: %s", im_func_name,
lua_typename(lua, t));
}
lsb_const_string output;
t = lua_type(lua, 1);
switch (t) {
case LUA_TUSERDATA:
{
heka_stream_reader *hsr = luaL_checkudata(lua, 1,
mozsvc_heka_stream_reader);
if (hsr->msg.raw.s) {
output.len = hsr->msg.raw.len;
output.s = hsr->msg.raw.s;
} else {
return luaL_error(lua, "%s() attempted to inject a nil message",
im_func_name);
}
}
break;
case LUA_TSTRING:
{
lsb_heka_message m;
lsb_init_heka_message(&m, 8);
output.s = lua_tolstring(lua, 1, &output.len);
bool ok = lsb_decode_heka_message(&m, output.s, output.len, NULL);
lsb_free_heka_message(&m);
if (!ok) {
return luaL_error(lua, "%s() attempted to inject a invalid protobuf "
"string", im_func_name);
}
}
break;
case LUA_TTABLE:
if (heka_encode_message_table(lsb, 1)) {
const char *err = lsb_get_error(lsb);
if (strlen(err) == 0) err = "exceeded output_limit";
return luaL_error(lua, "%s() failed: %s", im_func_name, err);
}
output.len = 0;
output.s = lsb_get_output(lsb, &output.len);
break;
default:
return luaL_error(lua, "%s() unsupported message type: %s",
im_func_name, lua_typename(lua, t));
}
lsb_heka_sandbox *hsb = lsb_get_parent(lsb);
if (hsb->cb.iim(hsb->parent, output.s, output.len, ncp, scp) != 0) {
return luaL_error(lua, "%s() failed: rejected by the callback",
im_func_name);
}
++hsb->stats.im_cnt;
hsb->stats.im_bytes += output.len;
return 0;
}
static int hj_parse_message(lua_State *lua)
{
lsb_lua_sandbox *lsb = static_cast<lsb_lua_sandbox *>
(lua_touserdata(lua, lua_upvalueindex(1)));
if (!lsb) {
return luaL_error(lua, "%s() invalid lightuserdata", __FUNCTION__);
}
int n = lua_gettop(lua);
int idx = 1;
lsb_heka_message *msg = NULL;
lsb_heka_sandbox *hsb = static_cast<lsb_heka_sandbox *>(lsb_get_parent(lsb));
if (hsb->type == 'i') {
luaL_argcheck(lua, n >= 2 && n <= 4, 0, "invalid number of arguments");
heka_stream_reader *hsr = static_cast<heka_stream_reader *>
(luaL_checkudata(lua, 1, mozsvc_heka_stream_reader));
msg = &hsr->msg;
idx = 2;
} else {
luaL_argcheck(lua, n >= 1 && n <= 3, 0, "invalid number of arguments");
if (!hsb->msg || !hsb->msg->raw.s) {
return luaL_error(lua, "no active message");
}
msg = hsb->msg;
}
lsb_const_string json = read_message(lua, idx, msg);
if (!json.s) return luaL_error(lua, "field not found");
char *inflated = NULL;
#ifdef HAVE_ZLIB
// automatically handle gzipped strings (optimization for Mozilla telemetry
// messages)
if (json.len > 2) {
if (json.s[0] == 0x1f && (unsigned char)json.s[1] == 0x8b) {
size_t mms = (size_t)lua_tointeger(lua, lua_upvalueindex(2));
inflated = lsb_ungzip(json.s, json.len, mms, NULL);
if (!inflated) return luaL_error(lua, "lsb_ungzip failed");
}
}
#endif
heka_json *hj = static_cast<heka_json *>(lua_newuserdata(lua, sizeof*hj));
hj->doc = new rj::Document;
hj->val = NULL;
hj->insitu = inflated;
hj->refs = new std::set<rj::Value *>;
luaL_getmetatable(lua, mozsvc_heka_json);
lua_setmetatable(lua, -2);
if (!hj->doc || !hj->refs) {
lua_pushstring(lua, "memory allocation failed");
return lua_error(lua);
}
bool err = false;
if (hj->insitu) {
if (hj->doc->ParseInsitu<rj::kParseStopWhenDoneFlag>(hj->insitu)
.HasParseError()) {
err = true;
lua_pushfstring(lua, "failed to parse offset:%f %s",
(lua_Number)hj->doc->GetErrorOffset(),
rj::GetParseError_En(hj->doc->GetParseError()));
}
} else {
rj::MemoryStream ms(json.s, json.len);
if (hj->doc->ParseStream<0, rj::UTF8<> >(ms).HasParseError()) {
err = true;
lua_pushfstring(lua, "failed to parse offset:%f %s",
(lua_Number)hj->doc->GetErrorOffset(),
rj::GetParseError_En(hj->doc->GetParseError()));
}
}
if (err) return lua_error(lua);
hj->refs->insert(hj->doc);
return 1;
}
