/*
 * 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 - stats allocator, minimal debug messages & no checks.
 * 5 - debug 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)
 * 0 - Only run GC when script hits it's memlimit.
 *
 * 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.
 *
 * FULL_GC -
 * 1 - Run full collection cycle.
 * 0 - Run small step collection.
 *
 * GC_STEP_SIZE_MUL -
 * Step size = "need size in kilobytes" * GC_STEP_SIZE_MUL
 *
 */
#define DEBUG_LEVEL 1
#define ZERO_FREE_MEM 0
#define ALWAYS_GC 0
#define ENABLE_GC 1
#define ALLOW_TEMP_ALLOCS 0
#define FULL_GC 0
#define GC_STEP_SIZE_MUL 1
#if 1
/* use stderr for debug messages. */
#define db_printf(fmt, args...) fprintf(stderr, fmt , ## args)
#else
#define db_printf(fmt, args...) printf(fmt , ## args)
#endif
#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;
	lua_Alloc frealloc;
	void			*ud;
	size_t		memused;
	size_t		peak_memused;
	size_t		max_memused;
	int				allow_tmp_allocs;
#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;
	size_t new_max_blocksize=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
			db_printf("block size=%d, count=%d\n", i, c);
#endif
			tot += i * c;
			new_max_blocksize=i;
		}
	}
	c = info->block_count[MAX_BLOCK_SIZE];
	if(c> 0) {
#if DEBUG_LEVEL>= 50
		db_printf("block size>= %d, count=%d\n", MAX_BLOCK_SIZE, c);
#endif
		tot += c;
		new_max_blocksize=MAX_BLOCK_SIZE;
	}
	info->max_blocksize=new_max_blocksize;
#if DEBUG_LEVEL>= 40
	db_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 {
			info->block_count[MAX_BLOCK_SIZE] -= osize;
			assert(info->block_count[MAX_BLOCK_SIZE]>= 0);
		}
	}
	if(nsize> 0) {
		if(nsize> info->max_blocksize) info->max_blocksize = nsize;
		if(nsize < MAX_BLOCK_SIZE) { assert(++(info->block_count[nsize])>= 0);
		} else {
			info->block_count[MAX_BLOCK_SIZE] += nsize;
			assert(info->block_count[MAX_BLOCK_SIZE]>= 0);
		}
	}
}
#endif
static void script_run_gc(script_info_t *info, size_t need)
{
#if FULL_GC
	(void)need;
	lua_gc(info->L, LUA_GCCOLLECT, 0);
#else
	size_t step_size = ((need>> 10) + 1) * GC_STEP_SIZE_MUL;
	int cycle_count = 0;
	do {
		if(lua_gc(info->L, LUA_GCSTEP, step_size)) {
			/* only allow completing the last cycle and starting a new cycle. */
			if((++cycle_count)> 1) break;
		}
	} while((info->memused + need)>= info->max_memused);
#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) {
		db_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) db_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
			db_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->allow_tmp_allocs == 0) {
#if DEBUG_LEVEL>= 40
		db_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->allow_tmp_allocs = ALLOW_TEMP_ALLOCS;
		script_run_gc(info, (nsize> osize)?(nsize - osize):0);
		info->allow_tmp_allocs = 0;
#if DEBUG_LEVEL>= 20
		info->gc_ptr = NULL;
#endif
#if DEBUG_LEVEL>= 40
		db_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;
			db_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
		db_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) db_printf("free 2 (%p), osize=%zd, bsize=%zd\n", ptr, osize, bsize);
#endif
#if DEBUG_LEVEL>= 20
	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
		db_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) db_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_stats_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(info->L != NULL && nsize> osize && info->memused>= info->max_memused && info->allow_tmp_allocs == 0) {
#if ENABLE_GC
		info->memused = old_size;
		/* try to free memory by collecting garbage. */
		info->allow_tmp_allocs = ALLOW_TEMP_ALLOCS;
		script_run_gc(info, (nsize> osize)?(nsize - osize):0);
		info->allow_tmp_allocs = 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
			db_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)
{
	size_t totalbytes = (lua_gc(L, LUA_GCCOUNT, 0) << 10) + lua_gc(L, LUA_GCCOUNTB, 0); size_t memlimit = (lua_gc(L, LUA_GCGETMEMLIMIT, 0) << 10); #if DEBUG_LEVEL>= 1
	script_info_t *info;
	lua_getallocf(L, (void *)(&info));
	db_printf("%s: memused=%zd, peak_memused=%zd\n", info->name, info->memused, info->peak_memused);
#endif
	printf("totalbytes=%zd, memlimit=%zd\n", totalbytes, memlimit);
	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);
#elif DEBUG_LEVEL>= 1
	L = lua_newstate(script_stats_alloc, info);
	lua_gc(L, LUA_GCSETMEMLIMIT, info->max_memused>> 10);
#else
	L = lua_open();
	lua_gc(L, LUA_GCSETMEMLIMIT, info->max_memused>> 10);
#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, int mem_limit)
{
	script_info_t *info;
	lua_State *L;
	int status;
	/* run script */
	info = (script_info_t *)calloc(1, sizeof(script_info_t));
	info->max_memused = mem_limit;
	info->allow_tmp_allocs = 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));
	}
	print_memstats(L);
	lua_close(L);
#if DEBUG_LEVEL>= 1
	db_printf("%s: memused=%zd, peak_memused=%zd\n", info->name,
		info->memused, info->peak_memused);
#endif
	/* check for memory leak. */
	assert(info->memused == 0);
	free(info);
	return 0;
}
int main(int argc, char *argv[])
{
	int mem_limit = 64 * 1024;
	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++) {
		if(argv[i][0] == '-') {
			switch(argv[i][1]) {
			case 'm':
				i++;
				if(i < argc) {
					mem_limit = atoi(argv[i]);
					printf("new mem_limit = %d\n", mem_limit);
				}
				break;
			default:
				printf("unkown option '%s'\n", argv[i]);
				break;
			}
			continue;
		}
		rc = run_script(argv[i], mem_limit);
	}
	return 0;
}
</div><div class="naked_ctrl">
<form action="/index.cgi/contrast" method="get" name="gate">
<p><a href="http://altstyle.alfasado.net">AltStyle</a> によって変換されたページ <a href="http://lua-users.org/files/wiki_insecure/users/RobertGabrielJakabosky/lua_memlimit.c">(-&gt;オリジナル)</a>
/ <label>アドレス: <input type="text" name="naked_post_url" value="http://lua-users.org/files/wiki_insecure/users/RobertGabrielJakabosky/lua_memlimit.c" size="22" /></label> <label>モード: <select name="naked_post_mode">
<option value="default">デフォルト</option>
<option value="speech">音声ブラウザ</option>
<option value="ruby">ルビ付き</option>
<option value="contrast" selected="selected">配色反転</option>
<option value="larger-text">文字拡大</option>
<option value="mobile">モバイル</option>
</select>
<input type="submit" value="表示" />
</p>
</form>
</div>