Handling of memory errors when creating to-be-closed upvalues

4 files changed, 148 insertions, 40 deletions

diff --git a/lfunc.c b/lfunc.c
index fde72b8c..4f9362f3 100644
--- a/lfunc.c
+++ b/lfunc.c
@@ -40,6 +40,7 @@ LClosure *luaF_newLclosure (lua_State *L, int n) {
   return c;
 }
 
+
 /*
 ** fill a closure with new closed upvalues
 */
@@ -56,31 +57,43 @@ void luaF_initupvals (lua_State *L, LClosure *cl) {
 }
 
 
+/*
+** Create a new upvalue with the given tag at the given level,
+** and link it to the list of open upvalues of 'L' after entry 'prev'.
+**/
+static UpVal *newupval (lua_State *L, int tag, StkId level, UpVal **prev) {
+  GCObject *o = luaC_newobj(L, tag, sizeof(UpVal));
+  UpVal *uv = gco2upv(o);
+  UpVal *next = *prev;
+  uv->v = s2v(level);  /* current value lives in the stack */
+  uv->u.open.next = next;  /* link it to list of open upvalues */
+  uv->u.open.previous = prev;
+  if (next)
+    next->u.open.previous = &uv->u.open.next;
+  *prev = uv;
+  if (!isintwups(L)) {  /* thread not in list of threads with upvalues? */
+    L->twups = G(L)->twups;  /* link it to the list */
+    G(L)->twups = L;
+  }
+  return uv;
+}
+
+
+/*
+** Find and reuse, or create if it does not exist, a regular upvalue
+** at the given level.
+*/
 UpVal *luaF_findupval (lua_State *L, StkId level) {
   UpVal **pp = &L->openupval;
-  GCObject *o;
   UpVal *p;
-  UpVal *uv;
   lua_assert(isintwups(L) || L->openupval == NULL);
-  while ((p = *pp) != NULL && uplevel(p) >= level) {
+  while ((p = *pp) != NULL && uplevel(p) >= level) {  /* search for it */
     if (uplevel(p) == level && !isdead(G(L), p))  /* corresponding upvalue? */
       return p;  /* return it */
     pp = &p->u.open.next;
   }
-  /* not found: create a new upvalue between 'pp' and 'p' */
-  o = luaC_newobj(L, LUA_TUPVAL, sizeof(UpVal));
-  uv = gco2upv(o);
-  uv->u.open.next = p;  /* link it to list of open upvalues */
-  uv->u.open.previous = pp;
-  if (p)
-    p->u.open.previous = &uv->u.open.next;
-  *pp = uv;
-  uv->v = s2v(level);  /* current value lives in the stack */
-  if (!isintwups(L)) {  /* thread not in list of threads with upvalues? */
-    L->twups = G(L)->twups;  /* link it to the list */
-    G(L)->twups = L;
-  }
-  return uv;
+  /* not found: create a new upvalue after 'pp' */
+  return newupval(L, LUA_TUPVAL, level, pp);
 }
 
 
@@ -89,25 +102,44 @@ static void callclose (lua_State *L, void *ud) {
 }
 
 
-static int closeupval (lua_State *L, UpVal *uv, StkId level, int status) {
-  StkId func = level + 1;  /* save slot for old error message */
-  if (status != LUA_OK)  /* was there an error? */
-    luaD_seterrorobj(L, status, level);  /* save error message */
-  else
-    setnilvalue(s2v(level));
-  if (ttisfunction(uv->v)) {  /* object to-be-closed is a function? */
-    setobj2s(L, func, uv->v);  /* will call it */
-    setobjs2s(L, func + 1, level);  /* error msg. as argument */
+/*
+** Prepare closing method with its argument for object at
+** index 'func' in the stack. Assume there is an error message
+** (or nil) just below the object.
+*/
+static int prepclosingmethod (lua_State *L, StkId func) {
+  if (ttisfunction(s2v(func))) {  /* object to-be-closed is a function? */
+    setobjs2s(L, func + 1, func - 1);  /* push error msg. as argument */
   }
   else {  /* try '__close' metamethod */
-    const TValue *tm = luaT_gettmbyobj(L, uv->v, TM_CLOSE);
-    if (ttisnil(tm))
-      return status;  /* no metamethod */
+    const TValue *tm = luaT_gettmbyobj(L, s2v(func), TM_CLOSE);
+    if (ttisnil(tm))  /* no metamethod? */
+      return 0;  /* nothing to call */
+    setobjs2s(L, func + 1, func);  /* 'self' is the argument */
     setobj2s(L, func, tm);  /* will call metamethod */
-    setobj2s(L, func + 1, uv->v);  /* with 'self' as argument */
   }
   L->top = func + 2;  /* add function and argument */
-  if (status == LUA_OK)  /* not in "error mode"? */
+  return 1;
+}
+
+
+/*
+** Prepare and call a closing method. If status is OK, code is
+** still inside the original protected call, and so any error
+** will be handled there. Otherwise, a previous error already
+** activated original protected call, and so the call to the
+** closing method must be protected here.
+*/
+static int closeupval (lua_State *L, TValue *uv, StkId level, int status) {
+  StkId func = level + 1;  /* save slot for old error message */
+  if (unlikely(status != LUA_OK))  /* was there an error? */
+    luaD_seterrorobj(L, status, level);  /* save error message */
+  else
+    setnilvalue(s2v(level));  /* no error message */
+  setobj2s(L, func, uv);  /* put object on top of error message */
+  if (!prepclosingmethod(L, func))
+    return status;  /* nothing to call */
+  if (likely(status == LUA_OK))  /* not in "error mode"? */
     callclose(L, func);  /* call closing method */
   else {  /* already inside error handler; cannot raise another error */
     int newstatus = luaD_pcall(L, callclose, func, savestack(L, level), 0);
@@ -118,6 +150,36 @@ static int closeupval (lua_State *L, UpVal *uv, StkId level, int status) {
 }
 
 
+/*
+** Try to create a to-be-closed upvalue
+** (can raise a memory-allocation error)
+*/
+static void trynewtbcupval (lua_State *L, void *ud) {
+  StkId level = cast(StkId, ud);
+  lua_assert(L->openupval == NULL || uplevel(L->openupval) < level);
+  newupval(L, LUA_TUPVALTBC, level, &L->openupval);
+}
+
+
+/*
+** Create a to-be-closed upvalue. If there is a memory error
+** when creating the upvalue, the closing method must be called here,
+** as there is no upvalue to call it later.
+*/
+void luaF_newtbcupval (lua_State *L, StkId level) {
+  int status = luaD_rawrunprotected(L, trynewtbcupval, level);
+  if (unlikely(status != LUA_OK)) {  /* memory error creating upvalue? */
+    StkId func = level + 1;
+    lua_assert(status == LUA_ERRMEM);
+    setobjs2s(L, func, level);  /* open space for error message */
+    luaD_seterrorobj(L, status, level);  /* save error message */
+    if (prepclosingmethod(L, func))
+      callclose(L, func);  /* call closing method */
+    luaD_throw(L, LUA_ERRMEM);  /* throw memory error */
+  }
+}
+
+
 void luaF_unlinkupval (UpVal *uv) {
   lua_assert(upisopen(uv));
   *uv->u.open.previous = uv->u.open.next;
@@ -139,7 +201,7 @@ int luaF_close (lua_State *L, StkId level, int status) {
     luaC_barrier(L, uv, slot);
     if (status >= 0 && uv->tt == LUA_TUPVALTBC) {  /* must be closed? */
       ptrdiff_t levelrel = savestack(L, level);
-      status = closeupval(L, uv, upl, status);  /* may reallocate the stack */
+      status = closeupval(L, uv->v, upl, status);  /* may reallocate the stack */
       level = restorestack(L, levelrel);
     }
   }
diff --git a/lfunc.h b/lfunc.h
index 4c788005..c9fe1314 100644
--- a/lfunc.h
+++ b/lfunc.h
@@ -47,6 +47,7 @@ LUAI_FUNC CClosure *luaF_newCclosure (lua_State *L, int nelems);
 LUAI_FUNC LClosure *luaF_newLclosure (lua_State *L, int nelems);
 LUAI_FUNC void luaF_initupvals (lua_State *L, LClosure *cl);
 LUAI_FUNC UpVal *luaF_findupval (lua_State *L, StkId level);
+LUAI_FUNC void luaF_newtbcupval (lua_State *L, StkId level);
 LUAI_FUNC int luaF_close (lua_State *L, StkId level, int status);
 LUAI_FUNC void luaF_unlinkupval (UpVal *uv);
 LUAI_FUNC void luaF_freeproto (lua_State *L, Proto *f);
diff --git a/lvm.c b/lvm.c
index e2994aa9..0d82756b 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1456,8 +1456,7 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_TBC) {
-        UpVal *up = luaF_findupval(L, ra);  /* create new upvalue */
-        up->tt = LUA_TUPVALTBC;  /* mark it to be closed */
+        luaF_newtbcupval(L, ra);  /* create new to-be-closed upvalue */
         vmbreak;
       }
       vmcase(OP_JMP) {
diff --git a/testes/locals.lua b/testes/locals.lua
index 8d55e9f5..d12c70a0 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -180,7 +180,7 @@ do
   do
     local scoped x = setmetatable({"x"}, {__close = function (self)
                                                    a[#a + 1] = self[1] end})
-    local scoped y = function () a[#a + 1] = "y" end
+    local scoped y = function (x) assert(x == nil); a[#a + 1] = "y" end
     a[#a + 1] = "in"
   end
   a[#a + 1] = "out"
@@ -210,27 +210,73 @@ do   -- errors in __close
          and #log == 4)
 end
 
-do
+if rawget(_G, "T") then
+  local function stack(n) n = (n == 0) or stack(n - 1); end;
   -- memory error inside closing function
   local function foo ()
-    local scoped y = function () io.write(2); T.alloccount() end
+    local scoped y = function () T.alloccount() end
     local scoped x = setmetatable({}, {__close = function ()
       T.alloccount(0); local x = {}   -- force a memory error
     end})
-    io.write("1\n")
     error("a")   -- common error inside the function's body
   end
 
+  stack(5)    -- ensure a minimal number of CI structures
+
+  -- despite memory error, 'y' will be executed and
+  -- memory limit will be lifted
   local _, msg = pcall(foo)
-T.alloccount()
   assert(msg == "not enough memory")
 
+  local function close (msg)
+    T.alloccount()
+    assert(msg == "not enough memory")
+  end
+
+  -- set a memory limit and return a closing function to remove the limit
+  local function enter (count)
+    stack(10)   -- reserve some stack space
+    T.alloccount(count)
+    return close
+  end
+
+  local function test ()
+    local scoped x = enter(0)   -- set a memory limit
+    -- creation of previous upvalue will raise a memory error
+    os.exit(false)    -- should not run
+  end
+
+  local _, msg = pcall(test)
+  assert(msg == "not enough memory")
+
+  -- now use metamethod for closing
+  close = setmetatable({}, {__close = function ()
+    T.alloccount()
+  end})
+
+  -- repeat test with extra closing upvalues
+  local function test ()
+    local scoped xxx = function (msg)
+      assert(msg == "not enough memory");
+      error(1000)   -- raise another error
+    end
+    local scoped xx = function (msg)
+      assert(msg == "not enough memory");
+    end
+    local scoped x = enter(0)   -- set a memory limit
+    -- creation of previous upvalue will raise a memory error
+    os.exit(false)    -- should not run
+  end
+
+  local _, msg = pcall(test)
+  assert(msg == 1000)
+
 end
 
 
 -- a suspended coroutine should not close its variables when collected
 local co = coroutine.wrap(function()
-  local scoped x = function () os.exit(1) end    -- should not run
+  local scoped x = function () os.exit(false) end    -- should not run
    coroutine.yield()
 end)
 co()