Fix handling of memory limit

qjs.c

@@ -96,6 +96,35 @@ static int eval_file(JSContext *ctx, const char *filename, int module)
     return ret;
 }
 
+static int64_t parse_limit(const char *arg) {
+    char *p;
+    unsigned long unit = 1024; /* default to traditional KB */
+    double d = strtod(arg, &p);
+
+    if (p == arg) {
+        fprintf(stderr, "Invalid limit: %s\n", arg);
+        return -1;
+    }
+
+    if (*p) {
+        switch (*p++) {
+        case 'b': case 'B': unit = 1UL <<  0; break;
+        case 'k': case 'K': unit = 1UL << 10; break; /* IEC kibibytes */
+        case 'm': case 'M': unit = 1UL << 20; break; /* IEC mebibytes */
+        case 'g': case 'G': unit = 1UL << 30; break; /* IEC gigibytes */
+        default:
+            fprintf(stderr, "Invalid limit: %s, unrecognized suffix, only k,m,g are allowed\n", arg);
+            return -1;
+        }
+        if (*p) {
+            fprintf(stderr, "Invalid limit: %s, only one suffix allowed\n", arg);
+            return -1;
+        }
+    }
+
+    return (int64_t)(d * unit);
+}
+
 static JSValue js_gc(JSContext *ctx, JSValue this_val,
                      int argc, JSValue *argv)
 {
@@ -199,7 +228,8 @@ static void *js_trace_malloc(JSMallocState *s, size_t size)
     /* Do not allocate zero bytes: behavior is platform dependent */
     assert(size != 0);
 
-    if (unlikely(s->malloc_size + size > s->malloc_limit))
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (unlikely(s->malloc_size + size > s->malloc_limit - 1))
         return NULL;
     ptr = malloc(size);
     js_trace_malloc_printf(s, "A %zd -> %p\n", size, ptr);
@@ -238,7 +268,8 @@ static void *js_trace_realloc(JSMallocState *s, void *ptr, size_t size)
         free(ptr);
         return NULL;
     }
-    if (s->malloc_size + size - old_size > s->malloc_limit)
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (s->malloc_size + size - old_size > s->malloc_limit - 1)
         return NULL;
 
     js_trace_malloc_printf(s, "R %zd %p", size, ptr);
@@ -273,8 +304,8 @@ void help(void)
            "    --std          make 'std' and 'os' available to the loaded script\n"
            "-T  --trace        trace memory allocation\n"
            "-d  --dump         dump the memory usage stats\n"
-           "    --memory-limit n       limit the memory usage to 'n' bytes\n"
-           "    --stack-size n         limit the stack size to 'n' bytes\n"
+           "    --memory-limit n       limit the memory usage to 'n' Kbytes\n"
+           "    --stack-size n         limit the stack size to 'n' Kbytes\n"
            "    --unhandled-rejection  dump unhandled promise rejections\n"
            "-q  --quit         just instantiate the interpreter and quit\n", JS_GetVersion());
     exit(1);
@@ -295,10 +326,10 @@ int main(int argc, char **argv)
     int module = -1;
     int load_std = 0;
     int dump_unhandled_promise_rejection = 0;
-    size_t memory_limit = 0;
     char *include_list[32];
     int i, include_count = 0;
-    size_t stack_size = 0;
+    int64_t memory_limit = -1;
+    int64_t stack_size = -1;
 
     argv0 = (JSCFunctionListEntry)JS_PROP_STRING_DEF("argv0", argv[0],
                                                      JS_PROP_C_W_E);
@@ -402,8 +433,7 @@ int main(int argc, char **argv)
                     }
                     opt_arg = argv[optind++];
                 }
-                // TODO(chqrlie): accept kmg suffixes
-                memory_limit = (size_t)strtod(opt_arg, NULL);
+                memory_limit = parse_limit(opt_arg);
                 break;
             }
             if (!strcmp(longopt, "stack-size")) {
@@ -414,8 +444,7 @@ int main(int argc, char **argv)
                     }
                     opt_arg = argv[optind++];
                 }
-                // TODO(chqrlie): accept kmg suffixes
-                stack_size = (size_t)strtod(opt_arg, NULL);
+                stack_size = parse_limit(opt_arg);
                 break;
             }
             if (opt) {
@@ -437,10 +466,10 @@ int main(int argc, char **argv)
         fprintf(stderr, "qjs: cannot allocate JS runtime\n");
         exit(2);
     }
-    if (memory_limit != 0)
-        JS_SetMemoryLimit(rt, memory_limit);
-    if (stack_size != 0)
-        JS_SetMaxStackSize(rt, stack_size);
+    if (memory_limit >= 0)
+        JS_SetMemoryLimit(rt, (size_t)memory_limit);
+    if (stack_size >= 0)
+        JS_SetMaxStackSize(rt, (size_t)stack_size);
     if (dump_flags != 0)
         JS_SetDumpFlags(rt, dump_flags);
     js_std_set_worker_new_context_func(JS_NewCustomContext);

quickjs.c

@@ -1608,7 +1608,7 @@ JSRuntime *JS_NewRuntime2(const JSMallocFunctions *mf, void *opaque)
 
     memset(&ms, 0, sizeof(ms));
     ms.opaque = opaque;
-    ms.malloc_limit = -1;
+    ms.malloc_limit = 0;
 
     rt = mf->js_malloc(&ms, sizeof(JSRuntime));
     if (!rt)
@@ -1685,7 +1685,8 @@ static void *js_def_malloc(JSMallocState *s, size_t size)
     /* Do not allocate zero bytes: behavior is platform dependent */
     assert(size != 0);
 
-    if (unlikely(s->malloc_size + size > s->malloc_limit))
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (unlikely(s->malloc_size + size > s->malloc_limit - 1))
         return NULL;
 
     ptr = malloc(size);
@@ -1723,7 +1724,8 @@ static void *js_def_realloc(JSMallocState *s, void *ptr, size_t size)
         free(ptr);
         return NULL;
     }
-    if (s->malloc_size + size - old_size > s->malloc_limit)
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (s->malloc_size + size - old_size > s->malloc_limit - 1)
         return NULL;
 
     ptr = realloc(ptr, size);

quickjs.h

@@ -291,6 +291,7 @@ typedef struct JSGCObjectHeader JSGCObjectHeader;
 JS_EXTERN JSRuntime *JS_NewRuntime(void);
 /* info lifetime must exceed that of rt */
 JS_EXTERN void JS_SetRuntimeInfo(JSRuntime *rt, const char *info);
+/* use 0 to disable memory limit */
 JS_EXTERN void JS_SetMemoryLimit(JSRuntime *rt, size_t limit);
 JS_EXTERN void JS_SetDumpFlags(JSRuntime *rt, uint64_t flags);
 JS_EXTERN void JS_SetGCThreshold(JSRuntime *rt, size_t gc_threshold);

v8.js

@@ -96,7 +96,7 @@ for (const file of files) {
     //print(`=== ${file}${envstr}${flagstr}`)
     print(`=== ${file}${envstr}`)
     const args = [argv0,
-                  "--stack-size", `${flags["--stack-size"]*1024}`,
+                  "--stack-size", `${flags["--stack-size"]}`,
                   "-I", "mjsunit.js",
                   "-I", tweak,
                   file]