implemented --hard-heap-limit and --heap-target-increment

Author: d-netto

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>]\"");
     }
 #ifdef _P64
     if (hint == 0) {

src/gc-stock.c

@@ -3212,7 +3212,7 @@ static int _jl_gc_collect(jl_ptls_t ptls, jl_gc_collection_t collection)
     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 == GC_HARD_HEAP_LIMIT_FOR_UNSET_FLAG) {
         // 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;
@@ -3293,6 +3293,23 @@ static int _jl_gc_collect(jl_ptls_t ptls, jl_gc_collection_t collection)
         target_heap = jl_atomic_load_relaxed(&gc_heap_stats.heap_target);
     }
 
+    // Kill the process if we are above the hard heap limit
+    if (heap_size > jl_options.hard_heap_limit) {
+        jl_errorf("Heap size of %zu bytes exceeds the hard limit of %zu bytes. "
+                  "Please increase the heap limit with `--hard-heap-limit`.",
+                  heap_size, jl_options.hard_heap_limit);
+    }
+    // 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 != GC_HARD_HEAP_LIMIT_FOR_UNSET_FLAG) {
+        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;
@@ -3690,6 +3707,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 != GC_HARD_HEAP_LIMIT_FOR_UNSET_FLAG) {
+        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;
@@ -3703,7 +3723,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>]\"");
     }
 #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)
 {
     long double value = 0.0;
     char unit[4] = {0};
@@ -151,6 +152,8 @@ JL_DLLEXPORT void jl_init_options(void)
                         0, // strip-ir
                         0, // permalloc_pkgimg
                         0, // heap-size-hint
+                        GC_HARD_HEAP_LIMIT_FOR_UNSET_FLAG, // hard-heap-limit
+                        GC_HEAP_TARGET_INCREMENT_FOR_UNSET_FLAG, // heap-target-increment
                         0, // trace_compile_timing
                         JL_TRIM_NO, // trim
                         0, // task_metrics
@@ -289,6 +292,16 @@ 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), T (tebibytes), or % (percentage\n"
+    "                                               of physical memory).\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), T (tebibytes),\n"
+    "                                               or % (percentage of physical memory).\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>]");
             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>]");
+            if (jl_options.hard_heap_limit == 0 || jl_options.hard_heap_limit == GC_HARD_HEAP_LIMIT_FOR_UNSET_FLAG)
+                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>]");
+            if (jl_options.heap_target_increment == 0 || jl_options.heap_target_increment == GC_HEAP_TARGET_INCREMENT_FOR_UNSET_FLAG)
+                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

@@ -2575,7 +2575,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);
 
 // Set julia-level ARGS array according to the arguments provided in
 // argc/argv

src/options.h

@@ -74,6 +74,9 @@
 // GC_TIME prints time taken by each phase of GC
 // #define GC_TIME
 
+#define GC_HARD_HEAP_LIMIT_FOR_UNSET_FLAG (UINT64_MAX)
+#define GC_HEAP_TARGET_INCREMENT_FOR_UNSET_FLAG (UINT64_MAX)
+
 // pool allocator configuration options
 
 // GC_SMALL_PAGE allocates objects in 4k pages

test/cmdlineargs.jl

@@ -1214,6 +1214,14 @@ 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()"`
+    exit_code = run(cmd)
+    @test exit_code != 0 # The process should be aborted due to the hard heap limit
+end
 end
 
 ## `Main.main` entrypoint
@@ -1253,6 +1261,38 @@ 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
+    @testset "--hard-heap-limit=$str" for (str, val) in [("1", 1), ("1e7", 1e7), ("2.5e7", 2.5e7), ("1MB", 1m), ("2.5g", 2.5g), ("1e4kB", 1e4k),
+        ("1e100", typemax(UInt64)), ("1e500g", typemax(UInt64)), ("1e-12t", 1), ("500000000b", 500000000)]
+        @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
+    @testset "--heap-target-increment=$str" for (str, val) in [("1", 1), ("1e7", 1e7), ("2.5e7", 2.5e7), ("1MB", 1m), ("2.5g", 2.5g), ("1e4kB", 1e4k),
+        ("1e100", typemax(UInt64)), ("1e500g", typemax(UInt64)), ("1e-12t", 1), ("500000000b", 500000000)]
+        @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