Add a notion of "free" objects during user thread allocation traversal

This way we can have objects that do not count towards the memory limit
for the specific script being run.

Author: HaroldCindy

VM/include/lua.h

@@ -7,6 +7,7 @@
 #include <stdint.h>
 #include <istream>
 #include <ostream>
+#include <unordered_set>
 
 #include "luaconf.h"
 
@@ -135,6 +136,9 @@ typedef int lua_Integer;
 // unsigned integer type
 typedef unsigned lua_Unsigned;
 
+// ServerLua: opaque set of GC object pointers for per-script memory accounting
+typedef std::unordered_set<void*> lua_OpaqueGCObjectSet;
+
 /*
 ** state manipulation
 */
@@ -396,7 +400,11 @@ LUA_API void lua_useconstsstate(lua_State* L, lua_State * constsL);
 // Get how much "real" memory is used by pages.
 LUA_API int lua_totalmemoverhead(lua_State *L);
 // Gets the total size of all user-allocated objects reachable from a user thread.
-LUA_API size_t lua_userthreadsize(lua_State *L);
+// If free_objects is provided, objects in that set will be excluded from the size calculation.
+LUA_API size_t lua_userthreadsize(lua_State *L, const lua_OpaqueGCObjectSet* free_objects);
+// Collects all memcat 2+ objects reachable from a user thread into a set.
+// Intended to be called immediately after luau_load() to capture bytecode constants.
+LUA_API lua_OpaqueGCObjectSet lua_collectfreeobjects(lua_State *L);
 
 
 /*

VM/src/lapi.cpp

@@ -2205,17 +2205,23 @@ CLANG_NOOPT void GCC_NOOPT lua_useconstsstate(lua_State *L, lua_State * constsL)
     luaC_validate(L);
 }
 
-size_t lua_userthreadsize(lua_State *L)
+size_t lua_userthreadsize(lua_State *L, const lua_OpaqueGCObjectSet* free_objects)
 {
     size_t total_size = 0;
     luaC_enumreachableuserallocs(
         L,
         &total_size,
         [](void* context, GCObject* ptr, uint8_t tt, uint8_t memcat, size_t size) {
             *((size_t*)context) += size;
-        }
+        },
+        free_objects
     );
     // Make sure we include the size of allocs from not-yet-rooted objects.
     // Note that this only makes sense if we call this on a thread that's currently executing.
     return total_size + L->global->unrooteduserallocs;
 }
+
+lua_OpaqueGCObjectSet lua_collectfreeobjects(lua_State *L)
+{
+    return luaC_collectfreeobjects(L);
+}

VM/src/lgc.h

@@ -143,11 +143,13 @@ LUAI_FUNC void luaC_enumheap(
     void (*node)(void* context, void* ptr, uint8_t tt, uint8_t memcat, size_t size, const char* name),
     void (*edge)(void* context, void* from, void* to, const char* name)
 );
-// ServerLua: User allocation tracking function
+// ServerLua: User allocation tracking functions
 LUAI_FUNC void luaC_enumreachableuserallocs(
     lua_State* L,
     void* context,
-    void (*node)(void* context, GCObject* ptr, uint8_t tt, uint8_t memcat, size_t size)
+    void (*node)(void* context, GCObject* ptr, uint8_t tt, uint8_t memcat, size_t size),
+    const lua_OpaqueGCObjectSet* free_objects
 );
+LUAI_FUNC lua_OpaqueGCObjectSet luaC_collectfreeobjects(lua_State* L);
 LUAI_FUNC int64_t luaC_allocationrate(lua_State* L);
 LUAI_FUNC const char* luaC_statename(int state);

VM/src/lgctraverse.cpp

@@ -22,7 +22,7 @@
 typedef struct ReachableContext
 {
     std::queue<GCObject*> queue;
-    std::unordered_set<GCObject*> visited;
+    std::unordered_set<void*> visited;
 } ReachableContext;
 
 static void enqueueobj(ReachableContext* ctx, GCObject* obj)
@@ -372,11 +372,18 @@ static size_t calcgcosize(GCObject *obj)
 void luaC_enumreachableuserallocs(
     lua_State* L,
     void* context,
-    void (*node)(void* context, GCObject* ptr, uint8_t tt, uint8_t memcat, size_t size)
+    void (*node)(void* context, GCObject* ptr, uint8_t tt, uint8_t memcat, size_t size),
+    const lua_OpaqueGCObjectSet* free_objects
 )
 {
     ReachableContext ctx;
 
+    // Pre-populate visited set with free objects if provided
+    if (free_objects)
+    {
+        ctx.visited = *free_objects;
+    }
+
     ctx.queue.push(obj2gco(L));
     ctx.visited.insert(obj2gco(L));
 
@@ -390,8 +397,30 @@ void luaC_enumreachableuserallocs(
             node(context, current, current->gch.tt, current->gch.memcat, calcgcosize(current));
 
         // Take any new references the current node has and add them to the queue
-        // Even if we don't want to include their size in the calculation, we may still want
-        // to traverse them.
         traverseobj(&ctx, current);
     }
 }
+
+lua_OpaqueGCObjectSet luaC_collectfreeobjects(lua_State* L)
+{
+    lua_OpaqueGCObjectSet free_objects;
+    ReachableContext ctx;
+
+    ctx.queue.push(obj2gco(L));
+    ctx.visited.insert(obj2gco(L));
+
+    while (!ctx.queue.empty())
+    {
+        GCObject* current = ctx.queue.front();
+        ctx.queue.pop();
+
+        // Collect memcat 2+ objects into the set
+        if (current->gch.memcat >= 2)
+            free_objects.insert(current);
+
+        // Traverse child references
+        traverseobj(&ctx, current);
+    }
+
+    return free_objects;
+}

VM/src/lll.cpp

@@ -806,7 +806,7 @@ static int ll_getfreememory(lua_State *L)
 {
     luaC_fullgc(L);
     luaX_graphheap(L, "/tmp/whatever.json");
-    luaSL_pushnativeinteger(L, (int)lua_userthreadsize(L));
+    luaSL_pushnativeinteger(L, (int)lua_userthreadsize(L, nullptr));
     return 1;
 }
 

tests/SLConformance.test.cpp

@@ -600,7 +600,7 @@ TEST_CASE("User thread alloc size calculation")
             // or the alloc would push us over the limit given the current approximate size.
             if (actual_size == 0 || (approximate_size + net_gain > MAX_MEM))
             {
-                approximate_size = actual_size = lua_userthreadsize(L);
+                approximate_size = actual_size = lua_userthreadsize(L, nullptr);
 
                 if (actual_size + net_gain > MAX_MEM)
                     return 1;