Lua: buffer system

balus

2021-11-14

Preface

最近在实现一套和 OpenResty 中的 ngx.re.* 类似的 API，其中 gsub/sub 这俩用到了 luaL_Buffer；在数据量比较小的情况下它工作的很好，可是一旦数据量上去了，程序就开始抛出异常，不过 OpenResty 却可以正常工作，着实困扰了我一阵；查阅资料发现网上说 luaL_Buffer 的确比较坑，怀疑是不是我用的方式不太对；最终通过看内部实现发现了问题所在，写这篇 blog 记录下

基本使用

可以简单将 luaL_Buffer 理解为一个 String Builder，即拼接字符串的工具，Lua 它的定义为：

A string buffer allows C code to build Lua strings piecemeal.

一般是这样使用的：

声明一个 luaL_Buffer 类型的变量 b
使用 luaL_Bufferinit(L, &b) 初始化它
通过 luaL_add*() 系列函数往该 buffer 中添加一部分字符串
通过调用 luaL_pushresult(&b) 结束字符串的构建，调用返回后会将最终得到的字符串放在栈顶

此外， manual 对 luaL_Buffer 还加了一个额外的描述：

During its normal operation, a string buffer uses a variable number of stack slots. So, while using a buffer, you cannot assume that you know where the top of the stack is. You can use the stack between successive calls to buffer operations as long as that use is balanced; that is, when you call a buffer operation, the stack is at the same level it was immediately after the previous buffer operation. (The only exception to this rule is luaL_addvalue After calling luaL_pushresult, the stack is back to its level when the buffer was initialized, plus the final string on its top.

大概是说，当操作 luaL_Buffer 时，它会使用一些栈空间，所以当使用 lua_Buffer 时，我们不能假定使用 luaL_Buffer 之前和之后栈没有变化；此外，在连续的 luaL_Buffer 操作之间，我们必须保持栈平衡

第一点很好理解，毕竟 buffer 想要动态地容纳任意长度的字符串，肯定是内部做了内存的分配，在栈上保存一些状态是必然的，至于用多少栈空间就取决于内部实现了。对于第二点，《Lua 程序设计》有做额外的解释：

…此外，尽管使用缓存区时我们可以将该栈用于其他用途，但是在访问栈(操作缓冲区)之前对栈的压入和弹出次数必须平衡

OpenResty 中是这样做的：

    ngx_str_t new_str;
    luaL_Buffer b;
    int count = 0;
    for ( ;; ) {
        count++;
        if (count == 1) {
            luaL_buffinit(L, &b);
        }

        /* call replace function to generate new string */
        lua_call(...);
        new_str.data = lua_tolstring(L, -1, &new_str.len);

        lua_insert(L, 1);

        luaL_addlstring(&b, &subj.data[cp_offset], cap[0] - cp_offset);
        luaL_addlstring(&b, new_str.data, new_str.len);

        lua_remove(L, 1);
    }

不用理会luaL_addlstring()添加的是什么内容，这里主要需要注意其中的 lua_insert() 和 lua_remove() 两处调用：

lua_insert() 将 replace function 产生的 new str 给推入栈底，这样就保证 buffer 在栈顶，同时也确保在后续的流程中 new str 不会被 GC
luaL_addlstring() 将 new str(的副本)拼接进去
lua_remove() 将 new str 从栈中移除

这样看似没有问题，而且实际上也工作得很好，但是一到 Lua5.4 中，大数据量的情况下它就无法工作；因为它实际上并没有保证栈是平衡的；在将 new str 推入栈底之后，栈的高度就发生了，buffer 所在的位置也变了，而这正是出错的根源。

源码剖析

源码之前，了无秘密；当然在看源码之前得看一下出错的栈帧，但是由于我还不知道怎么美观地贴图，所以就先不贴了。

结构定义

struct luaL_Buffer {
  char *b;  /* buffer address */
  size_t size;  /* buffer size */
  size_t n;  /* number of characters in buffer */
  lua_State *L;
  union {
    LUAI_MAXALIGN;  /* ensure maximum alignment for buffer */
    char b[LUAL_BUFFERSIZE];  /* initial buffer */
  } init;
};

有几点值得注意：

luaL_Buffer 内部有一个缓冲区(64bit 机器下其大小默认 1KB)，所以可以合理猜测当数据量比较小的时候，会直接存储在这个缓冲区中，而不需要动态分配内存；这和高版本的 std::string 的原理是类似的
luaL_Buffer 保存了 lua_State*，所以在它的 API 中才不需要像其他众多 API 一样以 lua_State* 作为第一个参数

buffer API

#define luaL_bufflen(bf)    ((bf)->n)
#define luaL_buffaddr(bf)   ((bf)->b)

#define luaL_addchar(B,c) \
  ((void)((B)->n < (B)->size || luaL_prepbuffsize((B), 1)), \
   ((B)->b[(B)->n++] = (c)))

#define luaL_addsize(B,s)   ((B)->n += (s))

#define luaL_buffsub(B,s)   ((B)->n -= (s))

LUALIB_API void (luaL_buffinit) (lua_State *L, luaL_Buffer *B);
LUALIB_API char *(luaL_prepbuffsize) (luaL_Buffer *B, size_t sz);
LUALIB_API void (luaL_addlstring) (luaL_Buffer *B, const char *s, size_t l);
LUALIB_API void (luaL_addstring) (luaL_Buffer *B, const char *s);
LUALIB_API void (luaL_addvalue) (luaL_Buffer *B);
LUALIB_API void (luaL_pushresult) (luaL_Buffer *B);
LUALIB_API void (luaL_pushresultsize) (luaL_Buffer *B, size_t sz);
LUALIB_API char *(luaL_buffinitsize) (lua_State *L, luaL_Buffer *B, size_t sz);

#define luaL_prepbuffer(B)  luaL_prepbuffsize(B, LUAL_BUFFERSIZE)

这些 API 大体可以分为 4 类：

初始化缓冲区
往缓冲区中追加各种类型的元素
获取缓冲区的最终结果
查询缓冲区的地址以及大小

初始化 buffer 结构体

Lua 提供了 luaL_buffinit() 和 luaL_buffinitsize() 两个函数：

当我们不知道结果会有多大时，就使用 luaL_buffinit()，然后 luaL_addXXX() 时让 Lua 动态分配内存
当我们知道结果的大小时，可以通过 luaL_bufferinitsize() 指定其大小，这个函数会返回缓冲区的空闲位置，然后我们手动将数据拷贝过去就 OK 了；这种情况下效率会更高

LUALIB_API void luaL_buffinit (lua_State *L, luaL_Buffer *B) {
  B->L = L;
  B->b = B->init.b;
  B->n = 0;
  B->size = LUAL_BUFFERSIZE;
  lua_pushlightuserdata(L, (void*)B);  /* push placeholder */
}

LUALIB_API char *luaL_buffinitsize (lua_State *L, luaL_Buffer *B, size_t sz) {
  luaL_buffinit(L, B);
  return prepbuffsize(B, sz, -1);
}

果不其然，和前面所说的一样，最开始的缓冲区就是用的 lua_Buffer 内部的 init 缓冲区。有一点需要注意，在初始化好了 luaL_Buffer 内部字段之后，Lua 将其作为一个 light userdata 推入了栈顶，此时它主要是作为一个占位符使用。

追加元素

这里以 luaL_addlstring 和 luaL_addvalue 为例，前者代表了绝大多数追加元素 API 的逻辑，后者则是一个特殊的 API：

LUALIB_API void luaL_addlstring (luaL_Buffer *B, const char *s, size_t l) {
  if (l > 0) {  /* avoid 'memcpy' when 's' can be NULL */
    char *b = prepbuffsize(B, l, -1);
    memcpy(b, s, l * sizeof(char));
    luaL_addsize(B, l);
  }
}

主要就 3 步：

获取可用的空闲缓冲区
将字符串拷贝至空闲缓冲区
更新缓冲区的已使用长度

luaL_addvalue 相比于其他 append API 的特殊之处在于，它是唯一个被调用时 box/buffer 不需要在栈顶的 API；因为我们在使用它时，需要先将元素推到栈顶，然后才进行拼接；此时 box/buffer 在 -2 这个位置，而其他的 append API 都假定 box/buffer 在栈顶(即 -1 位置)

LUALIB_API void luaL_addvalue (luaL_Buffer *B) {
  lua_State *L = B->L;
  size_t len;
  const char *s = lua_tolstring(L, -1, &len);
  char *b = prepbuffsize(B, len, -2);
  memcpy(b, s, len * sizeof(char));
  luaL_addsize(B, len);
  lua_pop(L, 1);  /* pop string */
}

获取空闲缓冲区

无论是哪个追加元素的 API，在真正将元素拷贝至缓冲区之前，都需要先确保缓存区中有足够的空间，而这就是 prebuffsize 的用途；

首先检查缓冲区剩余空间是否满足需要，如果满足则直接返回给调用方；

static char *prepbuffsize (luaL_Buffer *B, size_t sz, int boxidx) {
  checkbufferlevel(B, boxidx);
  if (B->size - B->n >= sz)  /* enough space? */
    return B->b + B->n;

不满足的话则进行后续操作，后面的逻辑涉及到 box 和 buffer 这俩概念，需要特殊说明一下。

  else {
    lua_State *L = B->L;
    char *newbuff;
    size_t newsize = newbuffsize(B, sz);
    /* create larger buffer */
    if (buffonstack(B))  /* buffer already has a box? */
      newbuff = (char *)resizebox(L, boxidx, newsize);  /* resize it */
    else {  /* no box yet */
      lua_remove(L, boxidx);  /* remove placeholder */
      newbox(L);  /* create a new box */
      lua_insert(L, boxidx);  /* move box to its intended position */
      lua_toclose(L, boxidx);
      newbuff = (char *)resizebox(L, boxidx, newsize);
      memcpy(newbuff, B->b, B->n * sizeof(char));  /* copy original content */
    }
    B->b = newbuff;
    B->size = newsize;
    return newbuff + B->n;
  }
}

此时需要重新分配一个更大的缓冲区，首先检查栈顶是不是 luaL_Buffer，如果是，那么说明缓冲区还是 luaL_Buffer 内部的；但是 luaL_Buffer 本身并不负责动态缓冲区的分配，负责这项工作的另有其人，在 buffer system 中被称为 UBox；此时需要将栈顶的 luaL_Buffer 结构替换为 UBox 结构，然后通过 resizebox() 分配更大的缓冲区。

UBox 结构以及相关函数定义如下：

typedef struct UBox {
  void *box;
  size_t bsize;
} UBox;

static void newbox (lua_State *L) {
  UBox *box = (UBox *)lua_newuserdatauv(L, sizeof(UBox), 0);
  box->box = NULL;
  box->bsize = 0;
  if (luaL_newmetatable(L, "_UBOX*"))  /* creating metatable? */
    luaL_setfuncs(L, boxmt, 0);  /* set its metamethods */
  lua_setmetatable(L, -2);
}

static void *resizebox (lua_State *L, int idx, size_t newsize) {
  void *ud;
  lua_Alloc allocf = lua_getallocf(L, &ud);
  UBox *box = (UBox *)lua_touserdata(L, idx);
  void *temp = allocf(ud, box->box, box->bsize, newsize);
  if (l_unlikely(temp == NULL && newsize > 0)) {  /* allocation error? */
    lua_pushliteral(L, "not enough memory");
    lua_error(L);  /* raise a memory error */
  }
  box->box = temp;
  box->bsize = newsize;
  return temp;
}

static int boxgc (lua_State *L) {
  resizebox(L, 1, 0);
  return 0;
}

static const luaL_Reg boxmt[] = {  /* box metamethods */
  {"__gc", boxgc},
  {"__close", boxgc},
  {NULL, NULL}
};

NOTE: 当 alloc 的第二个参数不为 NULL 时，lua_Alloc 相当于 realloc (第三个参数就是旧缓冲区的大小)，所以会将已有的数据拷贝到新的缓冲区去，我们无需额外操作；否则的话就相当于 malloc，具体的可以参考 lua_Alloc

为什么luaL_Buffer不负责内存的分配呢？这是因为在 buffer system 中，luaL_Buffer 对象由用户传入(一个指针)，这个对象并不归 Lua 管理，而是由 C 管理；而在 append API 中分配的内存对于 C 来说都是无感知的，需要由 Lua 管理；内存分配了自然就需要释放，但是我们并不知晓 C 会如何使用该 buffer (即我们不知道什么时候分配的内存可以释放了)，所以我们需要为它注册一个 gc handler，这样就可以保证内存会在适当的时机被释放。而 luaL_Buffer 只是一个 C 指针：一个 light userdata，我们无法(也不应该)为其注册元表与元方法，所以才需要换一个 UBox：一个 full userdata

而 UBox 也正是这样做的，只不过它在注册 __gc 元方法之外，还注册了一个 __close 元方法，这个是 Lua5.4 新引入的 TBC，实际上正是它导致的问题，这个后续再讲。

获取结果

LUALIB_API void luaL_pushresult (luaL_Buffer *B) {
  lua_State *L = B->L;
  checkbufferlevel(B, -1);
  lua_pushlstring(L, B->b, B->n);
  if (buffonstack(B))
    lua_closeslot(L, -2);  /* close the box */
  lua_remove(L, -2);  /* remove box or placeholder from the stack */
}

这里使用lua_pushlstring()将最终的字符串拷贝到栈顶，此时 UBox 就可以释放了，但是我们不能依赖这个逻辑，毕竟可能这个函数都没有被调用，所以还是得注册 gc handler，只是在 gc handler 中需要注意防止 double free

罪魁祸首：TBC

前面针对 OpenResty 中对 luaL_Buffer 的使用已经说了，它并没有保证栈的平衡，因为在两次 buffer 操作之间，buffer 的高度变了，这个在 Lua5.1 中并没有问题，但是 Lua5.4 中新增了 TBC(to be closed) 特性(具体的可以参考lua 5.4 可能会增加 to-be-closed 特性)，这个特性可以简单理解为 C++ 中的析构函数，Lua5.4 之前如果用 C 函数申请了一块资源，期望在使用完毕后可以清除干净，过去就只能依赖 __gc 方法。但 gc 的时机不可控，往往无法及时清理，而通过给变量增加 TBC 属性，当该变量超出其作用域时，就会执行其 __close 元方法，在其中我们可以销毁其资源

Lua5.4 中 tbc 变量是通过链表形式进行管理，当栈缩容，且被缩容的栈空间中含有 tbc，那么就需要将其销毁；

L->tbclist 记录着最后一个 tbc 节点，栈缩容时会判断该节点是否在缩容空间内，如果在，那么就根据这个节点调用缩容空间内所有 tbc 变量的 __close() 元方法；需要注意的是，这里节点的链接不是通过指针，而是通过相邻 tbc 变量在栈中的距离(在 Lua 的实现中算是一种比较常见的策略了)：

typedef union StackValue {
  TValue val;
  struct {
    TValuefields;
    unsigned short delta;
  } tbclist;
} StackValue;

所以判断 L->tbclist 是否在缩容空间内就很简单了：只是简单的做栈索引的对比；这也正是 Lua5.4 中无法继续使用 OpenResty 中的做法的原因。

下面是一个复现我所碰到的问题的 MVP：

int
main(int argc, char **argv)
{
  char dummy[1025];
  lua_State *L;
  luaL_Buffer b;

  L = luaL_newstate();
  luaL_openlibs(L);

  lua_pushliteral(L, "nothing special");
  luaL_buffinit(L, &b);
  luaL_addlstring(&b, dummy, 1025);
  lua_remove(L, 1);

  return 0;
}

首先推入一个短字符串
然后初始化 luaL_Buffer，此时 buffer 处在栈顶(idx = 2)
然后推入一个长度为 1025 的字符串(实际上只需两个字符串相加 > 1024 即可)，此时由于超出了 luaL_Buffer 内部缓存区的大小，所以会被转换为使用 UBox，并注册 __close() 方法将其声明为一个 tbc 变量，此时 L->tbclist = 2
然后调用 lua_remove()，其内部会调用 lua_rotate() 将 idx=1 和 idx=2 交换位置，将此时会调用 lua_settop(L, 1) 对栈进行缩容，此时 L->tbclist=2 在被缩容空间内，会尝试调用 index=2 的元素的 __close() 元方法，但是这个时候该位置上是一个 string，并没有注册该元方法，所以会报错：

PANIC: unprotected error in call to Lua API (attempt to call a nil value)