feat: add storage budget control for FileBackedMemory

Adds ability to limit FileBackedMemory usage via --storage_budget_gb flag or BB_STORAGE_BUDGET_GB env var. When budget is exceeded, falls back to RAM allocation.

Author: johnathan79717

barretenberg/cpp/src/barretenberg/api/api.hpp

@@ -23,6 +23,7 @@ class API {
         bool write_vk{ false };    // should we addditionally write the verification key when writing the proof
         bool include_gates_per_opcode{ false }; // should we include gates_per_opcode in the gates command output
         bool slow_low_memory{ false };          // use file backed memory for polynomials
+        size_t storage_budget_gb{ 0 };          // storage budget in GB for file backed memory (0 = no limit)
         bool update_inputs{ false };            // update inputs when check fails
 
         friend std::ostream& operator<<(std::ostream& os, const Flags& flags)
@@ -40,6 +41,7 @@ class API {
                << "  write_vk " << flags.write_vk << "\n"
                << "  include_gates_per_opcode " << flags.include_gates_per_opcode << "\n"
                << "  slow_low_memory " << flags.slow_low_memory << "\n"
+               << "  storage_budget_gb " << flags.storage_budget_gb << "\n"
                << "]" << std::endl;
             return os;
         }

barretenberg/cpp/src/barretenberg/bb/cli.cpp

@@ -278,6 +278,13 @@ int parse_and_run_cli_command(int argc, char* argv[])
             "--slow_low_memory", flags.slow_low_memory, "Enable low memory mode (can be 2x slower or more).");
     };
 
+    const auto add_storage_budget_option = [&](CLI::App* subcommand) {
+        return subcommand->add_option("--storage_budget_gb",
+                                      flags.storage_budget_gb,
+                                      "Storage budget in GB for FileBackedMemory. When exceeded, falls back to RAM "
+                                      "(requires --slow_low_memory).");
+    };
+
     const auto add_update_inputs_flag = [&](CLI::App* subcommand) {
         return subcommand->add_flag("--update_inputs", flags.update_inputs, "Update inputs if vk check fails.");
     };
@@ -351,6 +358,7 @@ int parse_and_run_cli_command(int argc, char* argv[])
     add_ipa_accumulation_flag(prove);
     remove_zk_option(prove);
     add_slow_low_memory_flag(prove);
+    add_storage_budget_option(prove);
     add_print_op_counts_flag(prove);
     add_op_counts_out_option(prove);
 
@@ -558,6 +566,9 @@ int parse_and_run_cli_command(int argc, char* argv[])
     debug_logging = flags.debug;
     verbose_logging = debug_logging || flags.verbose;
     slow_low_memory = flags.slow_low_memory;
+    if (flags.storage_budget_gb > 0) {
+        storage_budget = flags.storage_budget_gb * 1024ULL * 1024ULL * 1024ULL;
+    }
 #ifndef __wasm__
     if (print_op_counts || !op_counts_out.empty()) {
         bb::detail::use_op_count_time = true;

barretenberg/cpp/src/barretenberg/polynomials/backing_memory.cpp

@@ -1,5 +1,30 @@
 #include "barretenberg/polynomials/backing_memory.hpp"
+#include <atomic>
+#include <cstdlib>
+#include <limits>
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 bool slow_low_memory =
     std::getenv("BB_SLOW_LOW_MEMORY") == nullptr ? false : std::string(std::getenv("BB_SLOW_LOW_MEMORY")) == "1";
+
+// Parse storage budget from environment variable (in GB)
+static size_t parse_storage_budget()
+{
+    const char* env_val = std::getenv("BB_STORAGE_BUDGET_GB");
+    if (env_val == nullptr) {
+        return std::numeric_limits<size_t>::max(); // No limit by default
+    }
+
+    try {
+        size_t gb = std::stoull(env_val);
+        return gb * 1024ULL * 1024ULL * 1024ULL; // Convert GB to bytes
+    } catch (...) {
+        return std::numeric_limits<size_t>::max(); // Invalid value, no limit
+    }
+}
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+size_t storage_budget = parse_storage_budget();
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+std::atomic<size_t> current_storage_usage{ 0 };

barretenberg/cpp/src/barretenberg/polynomials/backing_memory.hpp

@@ -21,6 +21,12 @@
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 extern bool slow_low_memory;
 
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+extern size_t storage_budget;
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+extern std::atomic<size_t> current_storage_usage;
+
 template <typename T> class AlignedMemory;
 
 #ifndef __wasm__
@@ -43,7 +49,22 @@ template <typename Fr> class BackingMemory {
     {
 #ifndef __wasm__
         if (slow_low_memory) {
-            return std::shared_ptr<BackingMemory<Fr>>(new FileBackedMemory<Fr>(size));
+            size_t required_bytes = size * sizeof(Fr);
+            size_t current_usage = current_storage_usage.load();
+
+            // Check if we're under the storage budget
+            if (current_usage + required_bytes <= storage_budget) {
+                // Try to use FileBackedMemory
+                try {
+                    auto result = std::shared_ptr<BackingMemory<Fr>>(new FileBackedMemory<Fr>(size));
+                    current_storage_usage.fetch_add(required_bytes);
+                    return result;
+                } catch (const std::exception& e) {
+                    // FileBackedMemory allocation failed, fall back to AlignedMemory
+                    return std::shared_ptr<BackingMemory<Fr>>(new AlignedMemory<Fr>(size));
+                }
+            }
+            // Over budget, use AlignedMemory
         }
 #endif
         return std::shared_ptr<BackingMemory<Fr>>(new AlignedMemory<Fr>(size));
@@ -87,6 +108,8 @@ template <typename T> class FileBackedMemory : public BackingMemory<T> {
         }
         if (memory != nullptr && file_size > 0) {
             munmap(memory, file_size);
+            // Decrement storage usage when FileBackedMemory is freed
+            current_storage_usage.fetch_sub(file_size);
         }
         if (fd >= 0) {
             close(fd);