Re: PATCH: fixes bug with calling garbage collector from custom lua_Alloc

["Robert G. Jakabosky" <bobby@...>]

On Sunday 04, Robert G. Jakabosky wrote:
> memory limit.  This will allow shrinking the strings hashtable.  Attached
> is a new version of the lua_memlimit.c program that adds a
> 'allow_tmp_allocs' flag.
Sorry I attached an old copy.  The new version with the 'allow_tmp_allocs' 
flags is attached to this e-mail.

-- 
Robert G. Jakabosky

/*
 * Run Lua scripts with a memory limit.
 *
 * compile:
 * gcc -Wall -llua -o lua_memlimit lua_memlimit.c
 *
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#include <lua.h>
#include <lauxlib.h>
#include <lualib.h>

/*
 * DEBUG_LEVEL -
 * 1  - non-debug GC allocator, minimal debug messages & no checks.  See function 'script_alloc'
 *      it is a very basic memlimit + GC allocator.
 * 5  - debug GC allocator, with some checks enabled.
 * 10 - enables some more debug messages.
 * 20 - enables block size counter & block header.  Extra leak checking + block size validator.
 * 40 - enables a lot of debug messages.
 * 50 - enables all debug messages.
 *
 * ZERO_FREE_MEM -
 * 1 - fills all freed & allocated memory with zeros.
 *
 * ALWAYS_GC -
 * 1 - Do a full garbage collect on all allocations (where nsize > 0)
 *
 * ENABLE_GC -
 * 1 - Enables full garbage collection when max memory limit is hit.
 * 0 - No GC just returns NULL when max memory limit is hit.
 *
 * ALLOW_TEMP_ALLOCS -
 * 1 - Allow the script to temporary allocate more memory then it's limit during a full GC.
 * 0 - Use hardlimit and return NULL if script tries to pass it's limit during a full GC.
 */
#define DEBUG_LEVEL 5
#define ZERO_FREE_MEM 0
#define ALWAYS_GC 0
#define ENABLE_GC 1
#define ALLOW_TEMP_ALLOCS 1

#define MAX_BLOCK_SIZE 128 * 1024
#if DEBUG_LEVEL >= 20
#define BLOCK_HEADER	(sizeof(size_t)*2 + 16)
#else
#define BLOCK_HEADER	(0)
#endif
typedef struct {
	char			*name;
	lua_State *L;
	size_t		memused;
	size_t		peak_memused;
	size_t		max_memused;
	int				collecting;
#if DEBUG_LEVEL >= 20
	void			*gc_ptr;
	int				block_count[MAX_BLOCK_SIZE + 1];
	size_t		max_blocksize;
#endif
} script_info_t;

#if DEBUG_LEVEL >= 20
static void script_check_blocks(script_info_t *info)
{
	int i = 0;
	int c = 0;
	size_t tot = 0;

	for(i = 1; i < MAX_BLOCK_SIZE && i <= info->max_blocksize; i++) {
		c = info->block_count[i];
		if(c > 0) {
#if DEBUG_LEVEL >= 50
			printf("block size=%d, count=%d\n", i, c);
#endif
			tot += i * c;
		}
	}
	c = info->block_count[MAX_BLOCK_SIZE];
	if(c > 0) {
#if DEBUG_LEVEL >= 50
		printf("block size >= %d, count=%d\n", MAX_BLOCK_SIZE, c);
#endif
		tot += i * c;
	}
#if DEBUG_LEVEL >= 40
	printf("mem_total=%zd, max_blocksize=%zd\n", tot, info->max_blocksize);
#endif
	assert(tot == info->memused);
}

static void script_update_block_counts(script_info_t *info, size_t osize, size_t nsize)
{
	if(osize > 0) {
		if(osize < MAX_BLOCK_SIZE) {
			assert(--(info->block_count[osize]) >= 0);
		} else {
			assert(--(info->block_count[MAX_BLOCK_SIZE]) >= 0);
		}
	}
	if(nsize > 0) {
		if(nsize < MAX_BLOCK_SIZE) {
			assert(++(info->block_count[nsize]) >= 0);
		} else {
			assert(++(info->block_count[MAX_BLOCK_SIZE]) >= 0);
		}
	}
}
#endif

static void *script_debug_alloc(void *ud, void *ptr, size_t osize, size_t nsize)
{
	script_info_t *info=(script_info_t *)ud;
	size_t old_size = info->memused;
	int run_gc = 0;
#if DEBUG_LEVEL >= 20
	size_t bsize = -1;
#endif

#if DEBUG_LEVEL >= 40
	if(osize > info->memused) {
		printf("osize(%zd) > memused(%zd)\n", osize, info->memused);
	}
#endif
	assert(osize <= info->memused);

#if DEBUG_LEVEL >= 20
	if(ptr) {
		ptr -= BLOCK_HEADER;
		bsize = *((size_t *)ptr);
		assert(bsize == osize);
		assert(info->gc_ptr != ptr); /* check for second alloc call on same ptr. */
	}
#endif
	info->memused -= osize;
	if (nsize == 0) {
#if DEBUG_LEVEL >= 40
		if(ptr) printf("free 1 (%p), osize=%zd, bsize=%zd\n", ptr, osize, bsize);
#endif
#if DEBUG_LEVEL >= 20
		script_update_block_counts(info, osize, 0);
		script_check_blocks(info);
#endif
#if ZERO_FREE_MEM
		if(osize > 0) {
#if DEBUG_LEVEL >= 45
			printf("1 memset(%p,0,%zd)\n", ptr, osize + BLOCK_HEADER);
#endif
			memset(ptr, 0, osize + BLOCK_HEADER);
		}
#endif
		free(ptr);
		return NULL;
	}
	info->memused += nsize;
	if(nsize > osize && info->memused >= info->max_memused) {
#if ENABLE_GC
		run_gc = 1;
#if ALWAYS_GC
	} else if(info->L != NULL) {
		run_gc = 1;
#endif
	}
	if(run_gc && info->collecting == 0) {
#if DEBUG_LEVEL >= 40
		printf("LOW MEM: 1 osize=%zd, nsize=%zd, used=%zu, peak=%zu, need=%zd\n", osize, nsize,
			info->memused, info->peak_memused, (info->memused - info->max_memused));
#endif
		info->memused = old_size;
#if DEBUG_LEVEL >= 20
		info->gc_ptr = ptr;
#endif
		/* try to free memory by collecting garbage. */
		info->collecting = ALLOW_TEMP_ALLOCS;
		lua_gc(info->L, LUA_GCCOLLECT, 0);
		info->collecting = 0;
#if DEBUG_LEVEL >= 20
		info->gc_ptr = NULL;
#endif
#if DEBUG_LEVEL >= 40
		printf("LOW MEM: 2 used=%zu, peak=%zu\n", info->memused, info->peak_memused);
#endif
		/* check memory usage again. */
		old_size = info->memused;
		info->memused -= osize;
		info->memused += nsize;
		if(info->memused >= info->max_memused) {
			info->memused = old_size;
			printf("OUT OF MEMORY: memused=%zd, osize=%zd, nsize=%zd\n", info->memused, osize, nsize);
			return NULL;
		}
#else
		info->memused = old_size;
		return NULL;
#endif
	}
	if(info->memused > info->peak_memused) info->peak_memused = info->memused;
#if ZERO_FREE_MEM
	if(osize > nsize) {
#if DEBUG_LEVEL >= 45
		printf("2 memset(%p,0,%zd)\n", ptr + nsize + BLOCK_HEADER, osize - nsize);
#endif
		memset(ptr + nsize + BLOCK_HEADER, 0, osize - nsize);
	}
#endif
#if DEBUG_LEVEL >= 40
	if(ptr) printf("free 2 (%p), osize=%zd, bsize=%zd\n", ptr, osize, bsize);
#endif
#if DEBUG_LEVEL >= 20
	if(nsize > info->max_blocksize) info->max_blocksize = nsize;
	script_update_block_counts(info, osize, nsize);
	script_check_blocks(info);
#endif
	ptr = realloc(ptr, nsize + BLOCK_HEADER);
#if ZERO_FREE_MEM
	if(osize < nsize) {
#if DEBUG_LEVEL >= 45
		printf("3 memset(%p,0,%zd)\n", ptr + osize + BLOCK_HEADER, nsize - osize);
#endif
		memset(ptr + osize + BLOCK_HEADER, 0, nsize - osize);
	}
#endif
#if DEBUG_LEVEL >= 40
	if(ptr) printf("alloc  (%p), nsize=%zd\n", ptr, nsize);
#endif
#if DEBUG_LEVEL >= 20
	*((size_t *)ptr) = nsize;
	ptr += BLOCK_HEADER;
#endif
	return ptr;
}

static void *script_alloc(void *ud, void *ptr, size_t osize, size_t nsize)
{
	script_info_t *info=(script_info_t *)ud;
	size_t old_size = info->memused;

	info->memused -= osize;
	if (nsize == 0) {
		free(ptr);
		return NULL;
	}
	info->memused += nsize;
	if(nsize > osize && info->memused >= info->max_memused && info->collecting == 0) {
#if ENABLE_GC
		info->memused = old_size;
		/* try to free memory by collecting garbage. */
		info->collecting = ALLOW_TEMP_ALLOCS;
		lua_gc(info->L, LUA_GCCOLLECT, 0);
		info->collecting = 0;
		/* check memory usage again. */
		old_size = info->memused;
		info->memused -= osize;
		info->memused += nsize;
		if(info->memused >= info->max_memused) {
			info->memused = old_size;
#if DEBUG_LEVEL >= 1
			printf("OUT OF MEMORY: memused=%zd, osize=%zd, nsize=%zd\n", info->memused, osize, nsize);
#endif
			return NULL;
		}
#else
		info->memused = old_size;
		return NULL;
#endif
	}
	if(info->memused > info->peak_memused) info->peak_memused = info->memused;
	return realloc(ptr, nsize);
}

static int print_memstats(lua_State *L)
{
	script_info_t *info;

	lua_getallocf(L, (void *)(&info));

#if DEBUG_LEVEL >= 10
	printf("%s: before GC memused=%zd, peak_memused=%zd\n", info->name,
		info->memused, info->peak_memused);
	lua_gc(L, LUA_GCCOLLECT, 0);
	printf("%s: after GC  memused=%zd, peak_memused=%zd\n", info->name,
		info->memused, info->peak_memused);
#else
	printf("%s: memused=%zd, peak_memused=%zd\n", info->name,
		info->memused, info->peak_memused);
#endif
	return 0;
}

static int array_fill(lua_State *L)
{
	size_t l;
	const char *s = luaL_checklstring(L, 1, &l);
	int n = luaL_checkint(L, 2);
	int i;
	lua_createtable(L, n, 0);
	for(i=0; i < n; i++) {
		lua_pushstring(L, s);
		lua_rawseti(L, -2, i);
	}
	return 1;
}

static lua_State *create_newstate(script_info_t *info)
{
	lua_State *L;

#if DEBUG_LEVEL >= 5
	L = lua_newstate(script_debug_alloc, info);
#else
	L = lua_newstate(script_alloc, info);
#endif
	if(L == NULL) return L;
	/* load libs */
	info->L = L;
	luaL_openlibs(L);
	lua_register(L,"print_memstats", print_memstats);
	lua_register(L,"array_fill", array_fill);
	return L;
}

static int run_script(char *script_name)
{
	script_info_t *info;
	lua_State *L;
	int status;

	/* run script */
	info = (script_info_t *)calloc(1, sizeof(script_info_t));
	info->max_memused = 64 * 1024;
	info->collecting = 0;
	info->name = script_name;
	info->memused = 0;
	info->peak_memused = 0;
	/* create lua state & load script file. */
	L=create_newstate(info);
	status = luaL_loadfile(L, info->name);
	if(status != 0) {
		fprintf(stderr,"Failed to load script: %s\n",lua_tostring(L,-1));
		lua_close(L);
		free(info);
		return -1;
	}
	/* execute script. */
	status = lua_pcall(L, 0, LUA_MULTRET, 0);
	if(status != 0) {
		fprintf(stderr,"%s: %s\n", script_name,lua_tostring(L,-1));
	}
	lua_close(L);
	printf("%s: memused=%zd, peak_memused=%zd\n", info->name,
		info->memused, info->peak_memused);
	/* check for memory leak. */
	assert(info->memused == 0);
	free(info);
	return 0;
}

int main(int argc, char *argv[])
{
	int rc = 0;
	int i = 0;

	if(argc < 2) {
		printf("usage: %s <lua script 1> [<lua script 2>]\n", argv[0]);
		exit(1);
	}

	/* run scripts. */
	for(i=1; i < argc && rc == 0; i++) {
		rc = run_script(argv[i]);
	}

	return 0;
}