Introduce a few GC controls to limit the heap size when running benchmarks

base/options.jl

@@ -60,6 +60,8 @@ struct JLOptions
     strip_ir::Int8
     permalloc_pkgimg::Int8
     heap_size_hint::UInt64
+    hard_heap_limit::UInt64
+    heap_target_increment::UInt64
     trace_compile_timing::Int8
     trim::Int8
     task_metrics::Int8

src/gc-mmtk.c

@@ -78,7 +78,7 @@ void jl_gc_init(void) {
     if (jl_options.heap_size_hint == 0) {
         char *cp = getenv(HEAP_SIZE_HINT);
         if (cp)
-            hint = parse_heap_size_hint(cp, "JULIA_HEAP_SIZE_HINT=\"<size>[<unit>]\"");
+            hint = parse_heap_size_option(cp, "JULIA_HEAP_SIZE_HINT=\"<size>[<unit>]\"", 1);
     }
 #ifdef _P64
     if (hint == 0) {

src/gc-stock.c

@@ -3214,7 +3214,7 @@ static int _jl_gc_collect(jl_ptls_t ptls, jl_gc_collection_t collection) JL_NOTS
     uint64_t target_heap;
     const char *reason = ""; (void)reason; // for GC_TIME output stats
     old_heap_size = heap_size; // TODO: Update these values dynamically instead of just during the GC
-    if (collection == JL_GC_AUTO) {
+    if (collection == JL_GC_AUTO && jl_options.hard_heap_limit == 0) {
         // update any heuristics only when the user does not force the GC
         // but still update the timings, since GC was run and reset, even if it was too early
         uint64_t target_allocs = 0.0;
@@ -3295,6 +3295,27 @@ static int _jl_gc_collect(jl_ptls_t ptls, jl_gc_collection_t collection) JL_NOTS
         target_heap = jl_atomic_load_relaxed(&gc_heap_stats.heap_target);
     }
 
+    // Kill the process if we are above the hard heap limit
+    if (jl_options.hard_heap_limit != 0) {
+        if (heap_size > jl_options.hard_heap_limit) {
+            // Can't use `jl_errorf` here, because it will try to allocate memory
+            // and we are already at the hard limit.
+            jl_safe_printf("Heap size exceeded hard limit of %" PRIu64 " bytes.\n",
+                           jl_options.hard_heap_limit);
+            abort();
+        }
+    }
+    // Ignore heap limit computation from MemBalancer-like heuristics
+    // if the heap target increment goes above the value specified through
+    // `--heap-target-increment`.
+    // Note that if we reach this code, we can guarantee that the heap size
+    // is less than the hard limit, so there will be some room to grow the heap
+    // until the next GC without hitting the hard limit.
+    if (jl_options.heap_target_increment != 0) {
+        target_heap = heap_size + jl_options.heap_target_increment;
+        jl_atomic_store_relaxed(&gc_heap_stats.heap_target, target_heap);
+    }
+
     double old_ratio = (double)promoted_bytes/(double)heap_size;
     if (heap_size > user_max) {
         next_sweep_full = 1;
@@ -3692,6 +3713,9 @@ void jl_gc_init(void)
     arraylist_new(&finalizer_list_marked, 0);
     arraylist_new(&to_finalize, 0);
     jl_atomic_store_relaxed(&gc_heap_stats.heap_target, default_collect_interval);
+    if (jl_options.hard_heap_limit != 0) {
+        jl_atomic_store_relaxed(&gc_heap_stats.heap_target, jl_options.hard_heap_limit);
+    }
     gc_num.interval = default_collect_interval;
     gc_num.allocd = 0;
     gc_num.max_pause = 0;
@@ -3705,7 +3729,7 @@ void jl_gc_init(void)
     if (jl_options.heap_size_hint == 0) {
         char *cp = getenv(HEAP_SIZE_HINT);
         if (cp)
-            hint = parse_heap_size_hint(cp, "JULIA_HEAP_SIZE_HINT=\"<size>[<unit>]\"");
+            hint = parse_heap_size_option(cp, "JULIA_HEAP_SIZE_HINT=\"<size>[<unit>]\"", 1);
     }
 #ifdef _P64
     total_mem = uv_get_total_memory();

src/jloptions.c

@@ -3,6 +3,7 @@
 #include <limits.h>
 #include <errno.h>
 
+#include "options.h"
 #include "julia.h"
 #include "julia_internal.h"
 
@@ -36,7 +37,7 @@ JL_DLLEXPORT const char *jl_get_default_sysimg_path(void)
 
 /* This function is also used by gc-stock.c to parse the
  * JULIA_HEAP_SIZE_HINT environment variable. */
-uint64_t parse_heap_size_hint(const char *optarg, const char *option_name)
+uint64_t parse_heap_size_option(const char *optarg, const char *option_name, int allow_pct)
 {
     long double value = 0.0;
     char unit[4] = {0};
@@ -62,14 +63,16 @@ uint64_t parse_heap_size_hint(const char *optarg, const char *option_name)
             multiplier <<= 40;
             break;
         case '%':
-            if (value > 100)
-                jl_errorf("julia: invalid percentage specified in %s", option_name);
-            uint64_t mem = uv_get_total_memory();
-            uint64_t cmem = uv_get_constrained_memory();
-            if (cmem > 0 && cmem < mem)
-                mem = cmem;
-            multiplier = mem/100;
-            break;
+            if (allow_pct) {
+                if (value > 100)
+                    jl_errorf("julia: invalid percentage specified in %s", option_name);
+                uint64_t mem = uv_get_total_memory();
+                uint64_t cmem = uv_get_constrained_memory();
+                if (cmem > 0 && cmem < mem)
+                    mem = cmem;
+                multiplier = mem/100;
+                break;
+            }
         default:
             jl_errorf("julia: invalid argument to %s (%s)", option_name, optarg);
             break;
@@ -151,6 +154,8 @@ JL_DLLEXPORT void jl_init_options(void)
                         0, // strip-ir
                         0, // permalloc_pkgimg
                         0, // heap-size-hint
+                        0, // hard-heap-limit
+                        0, // heap-target-increment
                         0, // trace_compile_timing
                         JL_TRIM_NO, // trim
                         0, // task_metrics
@@ -289,6 +294,14 @@ static const char opts[]  =
     "                                               number of bytes, optionally in units of: B, K (kibibytes),\n"
     "                                               M (mebibytes), G (gibibytes), T (tebibytes), or % (percentage\n"
     "                                               of physical memory).\n\n"
+    " --hard-heap-limit=<size>[<unit>]              Set a hard limit on the heap size: if we ever go above this\n"
+    "                                               limit, we will abort. The value may be specified as a\n"
+    "                                               number of bytes, optionally in units of: B, K (kibibytes),\n"
+    "                                               M (mebibytes), G (gibibytes) or T (tebibytes).\n\n"
+    " --heap-target-increment=<size>[<unit>] Set an upper bound on how much the heap target\n"
+    "                                               can increase between consecutive collections. The value may be\n"
+    "                                               specified as a number of bytes, optionally in units of: B,\n"
+    "                                               K (kibibytes), M (mebibytes), G (gibibytes) or T (tebibytes).\n\n"
 ;
 
 static const char opts_hidden[]  =
@@ -380,6 +393,8 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
            opt_strip_metadata,
            opt_strip_ir,
            opt_heap_size_hint,
+           opt_hard_heap_limit,
+           opt_heap_target_increment,
            opt_gc_threads,
            opt_permalloc_pkgimg,
            opt_trim,
@@ -451,6 +466,8 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
         { "strip-ir",        no_argument,       0, opt_strip_ir },
         { "permalloc-pkgimg",required_argument, 0, opt_permalloc_pkgimg },
         { "heap-size-hint",  required_argument, 0, opt_heap_size_hint },
+        { "hard-heap-limit", required_argument, 0, opt_hard_heap_limit },
+        { "heap-target-increment", required_argument, 0, opt_heap_target_increment },
         { "trim",  optional_argument, 0, opt_trim },
         { 0, 0, 0, 0 }
     };
@@ -960,11 +977,23 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
             break;
         case opt_heap_size_hint:
             if (optarg != NULL)
-                jl_options.heap_size_hint = parse_heap_size_hint(optarg, "--heap-size-hint=<size>[<unit>]");
+                jl_options.heap_size_hint = parse_heap_size_option(optarg, "--heap-size-hint=<size>[<unit>]", 1);
             if (jl_options.heap_size_hint == 0)
                 jl_errorf("julia: invalid memory size specified in --heap-size-hint=<size>[<unit>]");
 
             break;
+        case opt_hard_heap_limit:
+            if (optarg != NULL)
+                jl_options.hard_heap_limit = parse_heap_size_option(optarg, "--hard-heap-limit=<size>[<unit>]", 0);
+            if (jl_options.hard_heap_limit == 0)
+                jl_errorf("julia: invalid memory size specified in --hard-heap-limit=<size>[<unit>]");
+            break;
+        case opt_heap_target_increment:
+            if (optarg != NULL)
+                jl_options.heap_target_increment = parse_heap_size_option(optarg, "--heap-target-increment=<size>[<unit>]", 0);
+            if (jl_options.heap_target_increment == 0)
+                jl_errorf("julia: invalid memory size specified in --heap-target-increment=<size>[<unit>]");
+            break;
         case opt_gc_threads:
             errno = 0;
             long nmarkthreads = strtol(optarg, &endptr, 10);

src/jloptions.h

@@ -64,6 +64,8 @@ typedef struct {
     int8_t strip_ir;
     int8_t permalloc_pkgimg;
     uint64_t heap_size_hint;
+    uint64_t hard_heap_limit;
+    uint64_t heap_target_increment;
     int8_t trace_compile_timing;
     int8_t trim;
     int8_t task_metrics;

src/julia.h

@@ -2580,7 +2580,7 @@ JL_DLLEXPORT ssize_t jl_sizeof_jl_options(void);
 JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp);
 JL_DLLEXPORT char *jl_format_filename(const char *output_pattern);
 
-uint64_t parse_heap_size_hint(const char *optarg, const char *option_name);
+uint64_t parse_heap_size_option(const char *optarg, const char *option_name, int allow_pct);
 
 // Set julia-level ARGS array according to the arguments provided in
 // argc/argv

test/cmdlineargs.jl

@@ -1214,6 +1214,16 @@ end
 
     @test readchomp(`$(Base.julia_cmd()) --startup-file=no --heap-size-hint=10M -e "println(@ccall jl_gc_get_max_memory()::UInt64)"`) == "$(1*1024*1024)"
 end
+
+@testset "hard heap limit" begin
+    # Set the hard heap limit to 100MB, try to allocate an array of 200MB
+    # and assert that the process is aborted, by checking the exit code.
+    cmd = `$(Base.julia_cmd()) --startup-file=no --hard-heap-limit=100M -e "a = Array{UInt8}(undef, 200*1024*1024); GC.gc()"`
+    p = open(pipeline(cmd, stderr=devnull, stdout=devnull))
+    exitcode = wait(p)
+    # The process should be aborted with an error code
+    @test exitcode != 0
+end
 end
 
 ## `Main.main` entrypoint
@@ -1253,6 +1263,72 @@ end
     end
 end
 
+@testset "--hard-heap-limit" begin
+    exename = `$(Base.julia_cmd())`
+    @test errors_not_signals(`$exename --hard-heap-limit -e "exit(0)"`)
+    @testset "--hard-heap-limit=$str" for str in ["asdf","","0","1.2vb","b","GB","2.5GB̂","1.2gb2","42gigabytes","5gig","2GiB","NaNt"]
+        @test errors_not_signals(`$exename --hard-heap-limit=$str -e "exit(0)"`)
+    end
+    k = 1024
+    m = 1024k
+    g = 1024m
+    t = 1024g
+    # Express one hundred megabytes as 100MB, 100m, 100e6, etc.
+    one_hundred_mb_strs_and_vals = [
+        ("100000000", 100000000), ("1e8", 1e8), ("100MB", 100m), ("100m", 100m), ("1e5kB", 1e5k),
+    ]
+    @testset "--hard-heap-limit=$str" for (str, val) in one_hundred_mb_strs_and_vals
+        @test parse(UInt64,read(`$exename --hard-heap-limit=$str -E "Base.JLOptions().hard_heap_limit"`, String)) == val
+    end
+    # Express two and a half gigabytes as 2.5g, 2.5GB, etc.
+    two_and_a_half_gigabytes_strs_and_vals = [
+        ("2500000000", 2500000000), ("2.5e9", 2.5e9), ("2.5g", 2.5g), ("2.5GB", 2.5g), ("2.5e6mB", 2.5e6m),
+    ]
+    @testset "--hard-heap-limit=$str" for (str, val) in two_and_a_half_gigabytes_strs_and_vals
+        @test parse(UInt64,read(`$exename --hard-heap-limit=$str -E "Base.JLOptions().hard_heap_limit"`, String)) == val
+    end
+    # Express one terabyte as 1TB, 1e12, etc.
+    one_terabyte_strs_and_vals = [
+        ("1000000000000", 1000000000000), ("1e12", 1e12), ("1TB", 1t), ("1e9gB", 1e9g),
+    ]
+    @testset "--hard-heap-limit=$str" for (str, val) in one_terabyte_strs_and_vals
+        @test parse(UInt64,read(`$exename --hard-heap-limit=$str -E "Base.JLOptions().hard_heap_limit"`, String)) == val
+    end
+end
+
+@testset "--heap-target-increment" begin
+    exename = `$(Base.julia_cmd())`
+    @test errors_not_signals(`$exename --heap-target-increment -e "exit(0)"`)
+    @testset "--heap-target-increment=$str" for str in ["asdf","","0","1.2vb","b","GB","2.5GB̂","1.2gb2","42gigabytes","5gig","2GiB","NaNt"]
+        @test errors_not_signals(`$exename --heap-target-increment=$str -e "exit(0)"`)
+    end
+    k = 1024
+    m = 1024k
+    g = 1024m
+    t = 1024g
+    # Express one hundred megabytes as 100MB, 100m, 100e6, etc.
+    one_hundred_mb_strs_and_vals = [
+        ("100000000", 100000000), ("1e8", 1e8), ("100MB", 100m), ("100m", 100m), ("1e5kB", 1e5k),
+    ]
+    @testset "--heap-target-increment=$str" for (str, val) in one_hundred_mb_strs_and_vals
+        @test parse(UInt64,read(`$exename --heap-target-increment=$str -E "Base.JLOptions().heap_target_increment"`, String)) == val
+    end
+    # Express two and a half gigabytes as 2.5g, 2.5GB, etc.
+    two_and_a_half_gigabytes_strs_and_vals = [
+        ("2500000000", 2500000000), ("2.5e9", 2.5e9), ("2.5g", 2.5g), ("2.5GB", 2.5g), ("2.5e6mB", 2.5e6m),
+    ]
+    @testset "--heap-target-increment=$str" for (str, val) in two_and_a_half_gigabytes_strs_and_vals
+        @test parse(UInt64,read(`$exename --heap-target-increment=$str -E "Base.JLOptions().heap_target_increment"`, String)) == val
+    end
+    # Express one terabyte as 1TB, 1e12, etc.
+    one_terabyte_strs_and_vals = [
+        ("1000000000000", 1000000000000), ("1e12", 1e12), ("1TB", 1t), ("1e9gB", 1e9g),
+    ]
+    @testset "--heap-target-increment=$str" for (str, val) in one_terabyte_strs_and_vals
+        @test parse(UInt64,read(`$exename --heap-target-increment=$str -E "Base.JLOptions().heap_target_increment"`, String)) == val
+    end
+end
+
 @testset "--timeout-for-safepoint-straggler" begin
     exename = `$(Base.julia_cmd())`
     timeout = 120