diff --git a/runtime/flan_rt.c b/runtime/flan_rt.c index 00a9ddef..f295fa44 100644 --- a/runtime/flan_rt.c +++ b/runtime/flan_rt.c @@ -1185,7 +1185,7 @@ static void *flan_heap_proc(flan_allocator *a, int32_t mode, void *p, * caller passes old_size for exactly this reason, and it is the one * number a wrong answer here would read off the end of. */ void *q; - if (flan_over_budget(a, size - old_size)) return NULL; + if (size > old_size && flan_over_budget(a, size - old_size)) return NULL; q = flan_heap_proc(a, FLAN_ALLOC_ALLOC, NULL, 0, size, align); if (!q) return NULL; if (p && old_size > 0) @@ -1331,6 +1331,9 @@ static void *flan_arena_proc(flan_allocator *a, int32_t mode, void *p, * without copying, which is the common shape. */ if (p && (uint8_t *)p + old_size == ar->base + ar->offset) { int64_t end = (int64_t)((uint8_t *)p - ar->base) + size; + /* The budget is checked here as on every other path; a block grown in + * place is still more live bytes. */ + if (size > old_size && flan_over_budget(a, size - old_size)) return NULL; if (end > ar->cap || end < 0) return NULL; ar->offset = end; if (end > ar->peak) ar->peak = end; diff --git a/test/programs/exhausted.flan b/test/programs/exhausted.flan index 206c99bc..0a1a1e2d 100644 --- a/test/programs/exhausted.flan +++ b/test/programs/exhausted.flan @@ -115,6 +115,24 @@ (println "INSERTIONSORT"))) ; INSERTIONSORT (println failures) ; 1 — failed once, retried once + ;; And an arena, whose budget is checked when a Vec grows its block in + ;; place as well as when it allocates a new one. A Vec that is the only thing + ;; pushing into an arena always grows in place, so without that check the + ;; ceiling would never be met. + (set tight (arena-new 65536)) + (set-alloc-budget tight 64) + (set failures 0) + (handler-bind + [(StorageExhausted [c] + (set failures (+ failures 1)) + (set-alloc-budget tight (* 2 (alloc-budget tight))) + (invoke-restart 'retry))] + (let [v (vec-new i32 tight)] + (dotimes [i 1000] (push v i)) + (println (length v)) ; 1000 + (println (at v 999)))) ; 999 + (println (> failures 0)) ; true + ;; And the restart is not once-per-program: it is established at each ;; allocation, so a later one offers it again. (set-alloc-budget tight 0) diff --git a/test/test_acceptance.ml b/test/test_acceptance.ml index 58a8cf4e..d8f8d00d 100644 --- a/test/test_acceptance.ml +++ b/test/test_acceptance.ml @@ -1663,7 +1663,7 @@ let () = flow an optimiser would otherwise launder. *) let exhausted_out = "64\n0\n126\ntrue\ntrue\n4\ntrue\n0\n8\n7\ntrue\n\ - 13\n73\n84\n90\nINSERTIONSORT\n1\n" + 13\n73\n84\n90\nINSERTIONSORT\n1\n1000\n999\ntrue\n" in outputs "storage exhausted, retried" "programs/exhausted.flan" exhausted_out; outputs ~opt:"-O0" "storage exhausted, retried, -O0" "programs/exhausted.flan"