Programming in Lua. Exercise 32.1

Lua.Exercise32.1.cpp

// Exercise 32.1: 
// Write a library that allows a script to limit the total amount
// of memory used by its Lua state.
// It may offer a single function, setlimit,
// to set that limit.

// The library should set its own allocation function.
// This function,
// before calling the original allocator,
// checks the total memory in use and returns NULL
// if the requested memory exceeds the limit.

// (Hint: the library can use lua_gc
// to initialize its byte count
// when it starts.
// It also can use the user data of the allocation function
// to keep its state:
// the byte count, the current memory limit, etc.;
// remember to use the original user data
// when calling the original allocation function.)

#define LUA_COMPAT_MODULE
#define LUA_COMPAT_APIINTCASTS

#define DEBUG_ALLOCATIONS false

#include <cstdint>
#include <iostream>
#include <pthread.h>

#include "lua.hpp"

namespace {

constexpr size_t kNoLimit = 0;
static void* kDefaultUserData = nullptr;
static lua_Alloc kDefaultAllocFunction = nullptr;

struct MemoryState {
    size_t used_memory;
    size_t memory_limit;
};

bool AbortWithMessage(const char* message) {
    std::cout << message << std::endl;
    std::abort();
}

bool LoadFile(lua_State* state, const std::string& filename) {
    return !(luaL_loadfile(state, filename.c_str()) || lua_pcall(state, 0, 0, 0));
}

int lua_memoryLib_setlimit(lua_State* L) {
    size_t memory_limit_in_bytes = static_cast<size_t>(luaL_checkinteger(L, 1));

    void* ud;
    lua_getallocf(L, &ud);
    MemoryState* memory_state_ptr = static_cast<MemoryState*>(ud);

    if (memory_state_ptr->memory_limit != kNoLimit) {
        luaL_argcheck(L, memory_limit_in_bytes <= memory_state_ptr->memory_limit, 1, "New memory limit is smaller than previously used limit which can lead to issues.");
    }

    memory_state_ptr->memory_limit = memory_limit_in_bytes;
    return 0;
}

int lua_memoryLib_getlimit(lua_State* L) {
    void* ud;
    lua_getallocf(L, &ud);
    MemoryState* memory_state_ptr = static_cast<MemoryState*>(ud);
    lua_pushnumber(L, memory_state_ptr->memory_limit);
    return 1;
}

int lua_memoryLib_getused(lua_State* L) {
    void* ud;
    lua_getallocf(L, &ud);
    MemoryState* memory_state_ptr = static_cast<MemoryState*>(ud);
    lua_pushnumber(L, memory_state_ptr->used_memory);
    return 1;
}

static const luaL_Reg lua_memoryLib_f[] = {
    {"setLimit", lua_memoryLib_setlimit},
    {"getLimit", lua_memoryLib_getlimit},
    {"getAllocatedBytes", lua_memoryLib_getused},
    {NULL, NULL}
};
      
int luaopen_memoryLib(lua_State* L) {
    luaL_newlib(L, lua_memoryLib_f);
    lua_setglobal(L, "memory");
    return 0;
}

void* lua_memoryLib_alloc(void *ud, void* ptr, size_t osize, size_t nsize) {
    MemoryState* memory_state_ptr = static_cast<MemoryState*>(ud);
    if (DEBUG_ALLOCATIONS) {
        printf("Custom allocation: old size %zu, new size %zu, used: %zu, limit: %zu, ptr: %p\n", osize, nsize, memory_state_ptr->used_memory,  memory_state_ptr->memory_limit, ptr); 
    }

    memory_state_ptr->used_memory -= osize;

    if (memory_state_ptr->memory_limit == kNoLimit ||
        (memory_state_ptr->used_memory + nsize <= memory_state_ptr->memory_limit)) {
        memory_state_ptr->used_memory += nsize;
    } else {
        printf("Memory size exceeded, avaliable: %zu, used: %zu, tried to allocated: %zu\n", memory_state_ptr->memory_limit, memory_state_ptr->used_memory, nsize); 
        std::abort();
    }

    return kDefaultAllocFunction(kDefaultUserData, ptr, osize, nsize);
}

size_t lua_getgcmemoryinuse(lua_State* L) {
    return (lua_gc(L, LUA_GCCOUNT, 0) * 1024)
        + lua_gc(L, LUA_GCCOUNTB, 0);
}

}  // namespace

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cout << "Please, specify lua script file to run." << std::endl;
        return 0;
    }

    lua_State* L = luaL_newstate();
    luaL_openlibs(L);
    luaopen_memoryLib(L);

    MemoryState memory_state;
    memory_state.used_memory = lua_getgcmemoryinuse(L);
    memory_state.memory_limit = kNoLimit;

    if (DEBUG_ALLOCATIONS) {
        printf("Init allocation. used: %zu, limit: %zu\n", memory_state.used_memory,  memory_state.memory_limit); 
    }

    kDefaultAllocFunction = lua_getallocf(L, &kDefaultUserData);
    lua_setallocf(L, lua_memoryLib_alloc, &memory_state);

    if (!LoadFile(L, std::string(argv[1]))) {
        AbortWithMessage(lua_tostring(L, -1));
        return -1;
    }

    return 0;
}

memory_analysis.lua

math.randomseed(os.time())

local charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"

function string.random(length)
  if length > 0 then
    return string.random(length - 1) .. charset:sub(math.random(1, #charset), 1)
  else
    return ""
  end
end

-- 5 Kb
-- memory.setLimit(5 * 1024)

buffer = {}

collectgarbage("collect")

print("Memory in bytes used by default:", memory.getAllocatedBytes(), "bytes")

print("Before adding elements:", memory.getAllocatedBytes(), "bytes")

for i=1,10 do
    buffer[i] = string.random(10)
    print("Memory in use:", memory.getAllocatedBytes(), "bytes")
end

print("After adding elements:", memory.getAllocatedBytes(), "bytes")

buffer = nil

collectgarbage("collect")

print("After running gc:", memory.getAllocatedBytes(), "bytes")

memory_analysis.txt

Memory in bytes used by default:	21736.0	bytes
Before adding elements:	21764.0	bytes
Memory in use:	23088.0	bytes
Memory in use:	23154.0	bytes
Memory in use:	23214.0	bytes
Memory in use:	23242.0	bytes
Memory in use:	23334.0	bytes
Memory in use:	23362.0	bytes
Memory in use:	23390.0	bytes
Memory in use:	23418.0	bytes
Memory in use:	23574.0	bytes
Memory in use:	23602.0	bytes
After adding elements:	23630.0	bytes
After running gc:	22584.0	bytes

pushing_the_limit.lua

math.randomseed(os.time())

local charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"

function string.random(length)
  if length > 0 then
    return string.random(length - 1) .. charset:sub(math.random(1, #charset), 1)
  else
    return ""
  end
end

-- 25 Kb
memory.setLimit(25 * 1024)

buffer = {}

for i=1,100 do
    buffer[i] = string.random(10)
    print("Memory in use:", memory.getAllocatedBytes(), "bytes")
end

pushing_the_limit.txt

Memory in use:	23399.0	bytes
Memory in use:	23443.0	bytes
Memory in use:	23548.0	bytes
Memory in use:	23576.0	bytes
Memory in use:	23668.0	bytes
Memory in use:	23696.0	bytes
Memory in use:	23724.0	bytes
Memory in use:	23752.0	bytes
Memory in use:	23908.0	bytes
Memory in use:	23936.0	bytes
Memory in use:	23964.0	bytes
Memory in use:	23992.0	bytes
Memory in use:	24020.0	bytes
Memory in use:	24048.0	bytes
Memory in use:	24076.0	bytes
Memory in use:	24104.0	bytes
Memory in use:	24388.0	bytes
Memory in use:	24416.0	bytes
Memory in use:	24444.0	bytes
Memory in use:	24472.0	bytes
Memory in use:	24500.0	bytes
Memory in use:	24528.0	bytes
Memory in use:	24556.0	bytes
Memory in use:	24584.0	bytes
Memory in use:	24612.0	bytes
Memory in use:	24640.0	bytes
Memory in use:	24668.0	bytes
Memory in use:	24696.0	bytes
Memory in use:	24724.0	bytes
Memory in use:	24752.0	bytes
Memory in use:	24780.0	bytes
Memory in use:	24808.0	bytes
Memory in use:	25348.0	bytes
Memory in use:	25376.0	bytes
Memory in use:	25404.0	bytes
Memory in use:	25432.0	bytes
Memory in use:	25460.0	bytes
Memory in use:	25488.0	bytes
Memory in use:	25516.0	bytes
Memory in use:	25544.0	bytes
Memory in use:Memory size exceeded, avaliable: 25600, used: 23524, tried to allocated: 4096
zsh: abort      ./lua-test ../main.lua