track heap usage

CMakeLists.txt

@@ -54,6 +54,7 @@ list(APPEND TEST_ENV "TEST_RUNTIME_EXE=${TOYWASM_CLI}")
 set(cli_sources
  "cli/main.c"
  "cli/repl.c"
+ "cli/str_to_uint.c"
 )
 
 add_executable(toywasm-cli ${cli_sources})

cli/main.c

@@ -7,7 +7,9 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "mem.h"
 #include "repl.h"
+#include "str_to_uint.h"
 #include "toywasm_config.h"
 #if defined(TOYWASM_ENABLE_WASI_THREADS)
 #include "wasi_threads.h"
@@ -34,6 +36,7 @@ enum longopt {
  opt_invoke,
  opt_load,
  opt_max_frames,
+ opt_max_memory,
  opt_max_stack_cells,
  opt_register,
  opt_repl,
@@ -125,6 +128,12 @@ static const struct option longopts[] = {
  NULL,
  opt_max_frames,
  },
+ {
+ "max-memory",
+ required_argument,
+ NULL,
+ opt_max_memory,
+ },
  {
  "max-stack-cells",
  required_argument,
@@ -226,6 +235,7 @@ static const char *opt_metavars[] = {
  [opt_repl_prompt] = "STRING",
  [opt_max_frames] = "NUMBER_OF_FRAMES",
  [opt_max_stack_cells] = "NUMBER_OF_CELLS",
+ [opt_max_memory] = "MEMORY_LIMIT_IN_BYTES",
 };
 
 static void
@@ -272,6 +282,11 @@ print_usage(void)
 int
 main(int argc, char *const *argv)
 {
+ struct mem_context mctx0, *mctx = &mctx0;
+ struct mem_context wasi_mctx0, *wasi_mctx = &wasi_mctx0;
+ struct mem_context dyld_mctx0, *dyld_mctx = &dyld_mctx0;
+ struct mem_context impobj_mctx0, *impobj_mctx = &impobj_mctx0;
+
  struct repl_state *state;
 #if defined(TOYWASM_ENABLE_WASI)
  VEC(, const char *) wasi_envs;
@@ -305,13 +320,25 @@ main(int argc, char *const *argv)
  xlog_tracing = 0;
 #endif
 #endif
+ mem_context_init(mctx);
+ mem_context_init(wasi_mctx);
+ mem_context_init(dyld_mctx);
+ mem_context_init(impobj_mctx);
+ wasi_mctx->parent = mctx;
+ dyld_mctx->parent = mctx;
+ impobj_mctx->parent = mctx;
 
  state = malloc(sizeof(*state));
  if (state == NULL) {
  goto fail;
  }
  toywasm_repl_state_init(state);
+ state->mctx = mctx;
+ state->wasi_mctx = wasi_mctx;
+ state->dyld_mctx = dyld_mctx;
+ state->impobj_mctx = impobj_mctx;
  struct repl_options *opts = &state->opts;
+ size_t limit;
  while ((ret = getopt_long(argc, argv, "", longopts, &longidx)) != -1) {
  switch (ret) {
  case opt_disable_jump_table:
@@ -329,7 +356,7 @@ main(int argc, char *const *argv)
  break;
 #if defined(TOYWASM_ENABLE_DYLD)
  case opt_dyld:
- if (state->nmodules > 0) {
+ if (state->modules.lsize > 0) {
  ret = EPROTO;
  goto fail;
  }
@@ -344,7 +371,7 @@ main(int argc, char *const *argv)
  break;
 #endif
  case opt_dyld_path:
- ret = VEC_PREALLOC(dyld_paths, 1);
+ ret = VEC_PREALLOC(mctx, dyld_paths, 1);
  if (ret != 0) {
  goto fail;
  }
@@ -353,7 +380,11 @@ main(int argc, char *const *argv)
  opts->dyld_options.paths = dyld_paths.p;
  break;
  case opt_dyld_stack_size:
- opts->dyld_options.stack_size = atoi(optarg);
+ ret = str_to_u32(optarg, 0,
+ &opts->dyld_options.stack_size);
+ if (ret != 0) {
+ goto fail;
+ }
  break;
 #endif
  case opt_invoke:
@@ -371,10 +402,31 @@ main(int argc, char *const *argv)
  might_need_help = false;
  break;
  case opt_max_frames:
- opts->exec_options.max_frames = atoi(optarg);
+ ret = str_to_u32(optarg, 0,
+ &opts->exec_options.max_frames);
+ if (ret != 0) {
+ goto fail;
+ }
+ break;
+ case opt_max_memory:
+ ret = str_to_size(optarg, 0, &limit);
+ if (ret != 0) {
+ goto fail;
+ }
+ ret = mem_context_setlimit(mctx, limit);
+ if (ret != 0) {
+ xlog_error("failed to set limit %zu (current "
+ "allocation: %zu)",
+ limit, mctx->allocated);
+ goto fail;
+ }
  break;
  case opt_max_stack_cells:
- opts->exec_options.max_stackcells = atoi(optarg);
+ ret = str_to_u32(optarg, 0,
+ &opts->exec_options.max_stackcells);
+ if (ret != 0) {
+ goto fail;
+ }
  break;
  case opt_register:
  ret = toywasm_repl_register(state, NULL, optarg);
@@ -420,7 +472,7 @@ main(int argc, char *const *argv)
  }
  break;
  case opt_wasi_env:
- ret = VEC_PREALLOC(wasi_envs, 1);
+ ret = VEC_PREALLOC(mctx, wasi_envs, 1);
  if (ret != 0) {
  goto fail;
  }
@@ -485,7 +537,7 @@ main(int argc, char *const *argv)
 #endif
  ) {
  const struct repl_module_state_u *mod_u =
- &state->modules[state->nmodules - 1];
+ &VEC_LASTELEM(state->modules);
  const struct repl_module_state *mod = &mod_u->u.repl;
  ret = wasi_threads_instance_set_thread_spawn_args(
  state->wasi_threads, mod->module, mod->extra_import);
@@ -538,12 +590,15 @@ main(int argc, char *const *argv)
 fail:
  toywasm_repl_reset(state);
 #if defined(TOYWASM_ENABLE_WASI)
- VEC_FREE(wasi_envs);
+ VEC_FREE(mctx, wasi_envs);
 #endif
 #if defined(TOYWASM_ENABLE_DYLD)
- VEC_FREE(dyld_paths);
+ VEC_FREE(mctx, dyld_paths);
 #endif
  free(state);
+ mem_context_clear(dyld_mctx);
+ mem_context_clear(wasi_mctx);
+ mem_context_clear(mctx);
  exit(exit_status);
 success:
  exit_status = 0;