executor: make task arena size configurable.

5 files changed, 275 insertions, 26 deletions

diff --git a/embassy-executor/Cargo.toml b/embassy-executor/Cargo.toml
index 3ff8440ca..ae46b17c6 100644
--- a/embassy-executor/Cargo.toml
+++ b/embassy-executor/Cargo.toml
@@ -27,6 +27,30 @@ default-target = "thumbv7em-none-eabi"
 targets = ["thumbv7em-none-eabi"]
 features = ["nightly", "defmt", "arch-cortex-m", "executor-thread", "executor-interrupt"]
 
+[dependencies]
+defmt = { version = "0.3", optional = true }
+log = { version = "0.4.14", optional = true }
+rtos-trace = { version = "0.1.2", optional = true }
+
+embassy-macros = { version = "0.2.1", path = "../embassy-macros" }
+embassy-time = { version = "0.1.5", path = "../embassy-time", optional = true}
+critical-section = "1.1"
+
+# needed for riscv
+# remove when https://github.com/rust-lang/rust/pull/114499 is merged
+portable-atomic = { version = "1.5", optional = true }
+
+# arch-cortex-m dependencies
+cortex-m = { version = "0.7.6", optional = true }
+
+# arch-wasm dependencies
+wasm-bindgen = { version = "0.2.82", optional = true }
+js-sys = { version = "0.3", optional = true }
+
+[dev-dependencies]
+critical-section = { version = "1.1", features = ["std"] }
+
+
 [features]
 
 # Architecture
@@ -49,25 +73,47 @@ turbowakers = []
 
 integrated-timers = ["dep:embassy-time"]
 
-[dependencies]
-defmt = { version = "0.3", optional = true }
-log = { version = "0.4.14", optional = true }
-rtos-trace = { version = "0.1.2", optional = true }
-
-embassy-macros = { version = "0.2.1", path = "../embassy-macros" }
-embassy-time = { version = "0.1.5", path = "../embassy-time", optional = true}
-critical-section = "1.1"
-
-# needed for riscv
-# remove when https://github.com/rust-lang/rust/pull/114499 is merged
-portable-atomic = { version = "1.5", optional = true }
-
-# arch-cortex-m dependencies
-cortex-m = { version = "0.7.6", optional = true }
+# BEGIN AUTOGENERATED CONFIG FEATURES
+# Generated by gen_config.py. DO NOT EDIT.
+task-arena-size-64 = []
+task-arena-size-128 = []
+task-arena-size-192 = []
+task-arena-size-256 = []
+task-arena-size-320 = []
+task-arena-size-384 = []
+task-arena-size-512 = []
+task-arena-size-640 = []
+task-arena-size-768 = []
+task-arena-size-1024 = []
+task-arena-size-1280 = []
+task-arena-size-1536 = []
+task-arena-size-2048 = []
+task-arena-size-2560 = []
+task-arena-size-3072 = []
+task-arena-size-4096 = [] # Default
+task-arena-size-5120 = []
+task-arena-size-6144 = []
+task-arena-size-8192 = []
+task-arena-size-10240 = []
+task-arena-size-12288 = []
+task-arena-size-16384 = []
+task-arena-size-20480 = []
+task-arena-size-24576 = []
+task-arena-size-32768 = []
+task-arena-size-40960 = []
+task-arena-size-49152 = []
+task-arena-size-65536 = []
+task-arena-size-81920 = []
+task-arena-size-98304 = []
+task-arena-size-131072 = []
+task-arena-size-163840 = []
+task-arena-size-196608 = []
+task-arena-size-262144 = []
+task-arena-size-327680 = []
+task-arena-size-393216 = []
+task-arena-size-524288 = []
+task-arena-size-655360 = []
+task-arena-size-786432 = []
+task-arena-size-1048576 = []
 
-# arch-wasm dependencies
-wasm-bindgen = { version = "0.2.82", optional = true }
-js-sys = { version = "0.3", optional = true }
-
-[dev-dependencies]
-critical-section = { version = "1.1", features = ["std"] }
+# END AUTOGENERATED CONFIG FEATURES
diff --git a/embassy-executor/README.md b/embassy-executor/README.md
index 3c1448a18..80ecfc71a 100644
--- a/embassy-executor/README.md
+++ b/embassy-executor/README.md
@@ -2,10 +2,36 @@
 
 An async/await executor designed for embedded usage.
 
-- No `alloc`, no heap needed. Task futures are statically allocated.
+- No `alloc`, no heap needed.
+- With nightly Rust, task futures can be fully statically allocated.
 - No "fixed capacity" data structures, executor works with 1 or 1000 tasks without needing config/tuning.
 - Integrated timer queue: sleeping is easy, just do `Timer::after_secs(1).await;`.
 - No busy-loop polling: CPU sleeps when there's no work to do, using interrupts or `WFE/SEV`.
 - Efficient polling: a wake will only poll the woken task, not all of them.
 - Fair: a task can't monopolize CPU time even if it's constantly being woken. All other tasks get a chance to run before a given task gets polled for the second time.
 - Creating multiple executor instances is supported, to run tasks with multiple priority levels. This allows higher-priority tasks to preempt lower-priority tasks.
+
+## Task arena
+
+When the `nightly` Cargo feature is not enabled, `embassy-executor` allocates tasks out of an arena (a very simple bump allocator).
+
+If the task arena gets full, the program will panic at runtime. To guarantee this doesn't happen, you must set the size to the sum of sizes of all tasks.
+
+Tasks are allocated from the arena when spawned for the first time. If the task exists, the allocation is not released to the arena, but can be reused to spawn the task again. For multiple-instance tasks (like `#[embassy_executor::task(pool_size = 4)]`), the first spawn will allocate memory for all instances. This is done for performance and to increase predictability (for example, spawning at least 1 instance of every task at boot guarantees an immediate panic if the arena is too small, while allocating instances on-demand could delay the panic to only when the program is under load).
+
+The arena size can be configured in two ways:
+
+- Via Cargo features: enable a Cargo feature like `task-arena-size-8192`. Only a selection of values
+  is available, check `Cargo.toml` for the list.
+- Via environment variables at build time: set the variable named `EMBASSY_EXECUTOR_TASK_ARENA_SIZE`. For example
+  `EMBASSY_EXECUTOR_TASK_ARENA_SIZE=4321 cargo build`. You can also set them in the `[env]` section of `.cargo/config.toml`.
+  Any value can be set, unlike with Cargo features.
+
+Environment variables take precedence over Cargo features. If two Cargo features are enabled for the same setting
+with different values, compilation fails.
+
+## Statically allocating tasks
+
+When using nightly Rust, enable the `nightly` Cargo feature. This will make `embassy-executor` use the `type_alias_impl_trait` feature to allocate all tasks in `static`s. Each task gets its own `static`, with the exact size to hold the task (or multiple instances of it, if using `pool_size`) calculated automatically at compile time. If tasks don't fit in RAM, this is detected at compile time by the linker. Runtime panics due to running out of memory are not possible.
+
+The configured arena size is ignored, no arena is used at all.
diff --git a/embassy-executor/build.rs b/embassy-executor/build.rs
index 6fe82b44f..07f31e3fb 100644
--- a/embassy-executor/build.rs
+++ b/embassy-executor/build.rs
@@ -1,6 +1,97 @@
-use std::env;
+use std::collections::HashMap;
+use std::fmt::Write;
+use std::path::PathBuf;
+use std::{env, fs};
+
+static CONFIGS: &[(&str, usize)] = &[
+    // BEGIN AUTOGENERATED CONFIG FEATURES
+    // Generated by gen_config.py. DO NOT EDIT.
+    ("TASK_ARENA_SIZE", 4096),
+    // END AUTOGENERATED CONFIG FEATURES
+];
+
+struct ConfigState {
+    value: usize,
+    seen_feature: bool,
+    seen_env: bool,
+}
 
 fn main() {
+    let crate_name = env::var("CARGO_PKG_NAME")
+        .unwrap()
+        .to_ascii_uppercase()
+        .replace('-', "_");
+
+    // only rebuild if build.rs changed. Otherwise Cargo will rebuild if any
+    // other file changed.
+    println!("cargo:rerun-if-changed=build.rs");
+
+    // Rebuild if config envvar changed.
+    for (name, _) in CONFIGS {
+        println!("cargo:rerun-if-env-changed={crate_name}_{name}");
+    }
+
+    let mut configs = HashMap::new();
+    for (name, default) in CONFIGS {
+        configs.insert(
+            *name,
+            ConfigState {
+                value: *default,
+                seen_env: false,
+                seen_feature: false,
+            },
+        );
+    }
+
+    let prefix = format!("{crate_name}_");
+    for (var, value) in env::vars() {
+        if let Some(name) = var.strip_prefix(&prefix) {
+            let Some(cfg) = configs.get_mut(name) else {
+                panic!("Unknown env var {name}")
+            };
+
+            let Ok(value) = value.parse::<usize>() else {
+                panic!("Invalid value for env var {name}: {value}")
+            };
+
+            cfg.value = value;
+            cfg.seen_env = true;
+        }
+
+        if let Some(feature) = var.strip_prefix("CARGO_FEATURE_") {
+            if let Some(i) = feature.rfind('_') {
+                let name = &feature[..i];
+                let value = &feature[i + 1..];
+                if let Some(cfg) = configs.get_mut(name) {
+                    let Ok(value) = value.parse::<usize>() else {
+                        panic!("Invalid value for feature {name}: {value}")
+                    };
+
+                    // envvars take priority.
+                    if !cfg.seen_env {
+                        if cfg.seen_feature {
+                            panic!("multiple values set for feature {}: {} and {}", name, cfg.value, value);
+                        }
+
+                        cfg.value = value;
+                        cfg.seen_feature = true;
+                    }
+                }
+            }
+        }
+    }
+
+    let mut data = String::new();
+
+    for (name, cfg) in &configs {
+        writeln!(&mut data, "pub const {}: usize = {};", name, cfg.value).unwrap();
+    }
+
+    let out_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());
+    let out_file = out_dir.join("config.rs").to_string_lossy().to_string();
+    fs::write(out_file, data).unwrap();
+
+    // cortex-m targets
     let target = env::var("TARGET").unwrap();
 
     if target.starts_with("thumbv6m-") {
diff --git a/embassy-executor/gen_config.py b/embassy-executor/gen_config.py
new file mode 100644
index 000000000..e427d29f4
--- /dev/null
+++ b/embassy-executor/gen_config.py
@@ -0,0 +1,82 @@
+import os
+
+abspath = os.path.abspath(__file__)
+dname = os.path.dirname(abspath)
+os.chdir(dname)
+
+features = []
+
+
+def feature(name, default, min=None, max=None, pow2=None, vals=None, factors=[]):
+    if vals is None:
+        assert min is not None
+        assert max is not None
+
+        vals = set()
+        val = min
+        while val <= max:
+            vals.add(val)
+            for f in factors:
+                if val * f <= max:
+                    vals.add(val * f)
+            if (pow2 == True or (isinstance(pow2, int) and val >= pow2)) and val > 0:
+                val *= 2
+            else:
+                val += 1
+        vals.add(default)
+        vals = sorted(list(vals))
+
+    features.append(
+        {
+            "name": name,
+            "default": default,
+            "vals": vals,
+        }
+    )
+
+
+feature(
+    "task_arena_size", default=4096, min=64, max=1024 * 1024, pow2=True, factors=[3, 5]
+)
+
+# ========= Update Cargo.toml
+
+things = ""
+for f in features:
+    name = f["name"].replace("_", "-")
+    for val in f["vals"]:
+        things += f"{name}-{val} = []"
+        if val == f["default"]:
+            things += " # Default"
+        things += "\n"
+    things += "\n"
+
+SEPARATOR_START = "# BEGIN AUTOGENERATED CONFIG FEATURES\n"
+SEPARATOR_END = "# END AUTOGENERATED CONFIG FEATURES\n"
+HELP = "# Generated by gen_config.py. DO NOT EDIT.\n"
+with open("Cargo.toml", "r") as f:
+    data = f.read()
+before, data = data.split(SEPARATOR_START, maxsplit=1)
+_, after = data.split(SEPARATOR_END, maxsplit=1)
+data = before + SEPARATOR_START + HELP + things + SEPARATOR_END + after
+with open("Cargo.toml", "w") as f:
+    f.write(data)
+
+
+# ========= Update build.rs
+
+things = ""
+for f in features:
+    name = f["name"].upper()
+    things += f'    ("{name}", {f["default"]}),\n'
+
+SEPARATOR_START = "// BEGIN AUTOGENERATED CONFIG FEATURES\n"
+SEPARATOR_END = "// END AUTOGENERATED CONFIG FEATURES\n"
+HELP = "    // Generated by gen_config.py. DO NOT EDIT.\n"
+with open("build.rs", "r") as f:
+    data = f.read()
+before, data = data.split(SEPARATOR_START, maxsplit=1)
+_, after = data.split(SEPARATOR_END, maxsplit=1)
+data = before + SEPARATOR_START + HELP + things + "    " + SEPARATOR_END + after
+with open("build.rs", "w") as f:
+    f.write(data)
diff --git a/embassy-executor/src/lib.rs b/embassy-executor/src/lib.rs
index ac7dbb035..d8ac4893b 100644
--- a/embassy-executor/src/lib.rs
+++ b/embassy-executor/src/lib.rs
@@ -40,6 +40,11 @@ pub mod raw;
 mod spawner;
 pub use spawner::*;
 
+mod config {
+    #![allow(unused)]
+    include!(concat!(env!("OUT_DIR"), "/config.rs"));
+}
+
 /// Implementation details for embassy macros.
 /// Do not use. Used for macros and HALs only. Not covered by semver guarantees.
 #[doc(hidden)]
@@ -86,7 +91,7 @@ pub mod _export {
             let align_offset = (ptr as usize).next_multiple_of(layout.align()) - (ptr as usize);
 
             if align_offset + layout.size() > bytes_left {
-                panic!("arena full");
+                panic!("embassy-executor: task arena is full. You must increase the arena size, see the documentation for details: https://docs.embassy.dev/embassy-executor/");
             }
 
             let res = unsafe { ptr.add(align_offset) };
@@ -98,8 +103,7 @@ pub mod _export {
         }
     }
 
-    const ARENA_SIZE: usize = 16 * 1024;
-    static ARENA: Arena<ARENA_SIZE> = Arena::new();
+    static ARENA: Arena<{ crate::config::TASK_ARENA_SIZE }> = Arena::new();
 
     pub struct TaskPoolRef {
         // type-erased `&'static mut TaskPool<F, N>`