LuaJITのメモリアロケータを入れ替える（気をつけないとメモリリークする）

luajit lua mysql

環境

  • Lua 5.1.4
  • LuaJIT 2.0.0-beta9

動機

MySQLのUDFの中でLuaを実行したい。
問題は2つ。

  • 1コネクションにつきLuaインスタンスを1つ作るので、max_connectionsで設定した数だけLuaインスタンスが同時に立ち上がる可能性がある。
  • 1つメモリを食いつぶすクエリがあるだけで、他のすべてのクエリまで影響を受けるのは避けたい。

ということで、Luaインスタンス毎にメモリ使用量の上限を設定したい。

実装

#include "lua.h"
#include "lauxlib.h"
#include "lualib.h"
#include "luajit.h"

typedef struct st_mydata {
  lua_Alloc old_allocf;
  void *old_allocd;
  size_t memory_usage;
  size_t memory_limit;
} MYDATA;

void *myallocf(void *ud, void *ptr, size_t osize, size_t nsize) {
  MYDATA *mydata = (MYDATA *)ud;
  mydata->memory_usage += nsize - osize;
  if (nsize != 0 &&
      nsize > osize &&
      mydata->memory_usage > mydata->memory_limit) {
    return NULL;
  }
  return mydata->old_allocf(mydata->old_allocd, ptr, osize, nsize);
}

int main() {
  MYDATA mydata = {0};
  lua_State *lua = luaL_newstate();
  if (lua); else return 1;
  mydata.old_allocf = lua_getallocf(lua, &mydata.old_allocd);
  lua_setallocf(lua, myallocf, &mydata);
  // ...
  lua_close(lua);
  return 0;
}

myallocf関数が、自前で実装した、容量制限付きのメモリアロケータ。
仕組みは単純で、制限容量を超える場合にはNULLを返し、超えない場合には処理系組み込みのアロケータを呼び出すだけ。

Lua 5.1.4では正しく動作するけど、LuaJIT 2.0.0-beta9ではメモリリークする。
便宜上main関数になってるけど、実際にはUDFの中(=MySQLの1つのスレッドの中)でLuaインスタンスを生成・削除するので、メモリリークが問題になる。

原因

lj_state.cの155行目

static void close_state(lua_State *L)
{
  // ...
#ifndef LUAJIT_USE_SYSMALLOC
  if (g->allocf == lj_alloc_f)
    lj_alloc_destroy(g->allocd);
  else
#endif
    g->allocf(g->allocd, G2GG(g), sizeof(GG_State), 0);
}

allocf(とallocd)を差し替えたままlua_closeすると、元のallocdが開放されないので、メモリリークになる。
単純に、lua_closeする前に、元のallocfに戻しておけばいい。
getallocfやsetallocfは例外を投げない。

int main() {
  MYDATA mydata = {0};
  lua_State *lua = luaL_newstate();
  if (lua); else return 1;
  mydata.old_allocf = lua_getallocf(lua, &mydata.old_allocd);
  lua_setallocf(lua, myallocf, &mydata);
  // ...
  lua_setallocf(lua, mydata.old_allocf, mydata.old_allocd); // 追加
  lua_close(lua);
  return 0;
}

最後に

MySQLに対して、Luaのコードを投げてHandler APIを叩けるUDFを作ってます。

クエリのサンプルはこんな感じ。

SELECT mylua('
  local t = {};
  mylua.init_table("localuser", "timeline", "user_id", "user_id", "status_id");
  mylua.index_read_map(mylua.HA_READ_KEY_OR_NEXT, mylua.arg.uid, mylua.arg.sid);
  table.insert(t, { mylua.val_int("user_id"), mylua.val_int("status_id") });
  mylua.index_next();
  table.insert(t, { mylua.val_int("user_id"), mylua.val_int("status_id") });
  return t;
', '{"uid":3,"sid":32}');

まだ本番では使っていませんが、だいぶ安定してきたので、興味がある方はどうぞ。githubで公開しています。
https://github.com/atsumu/mylua