[web_server_idf] Build state documents in a stack arena so a state event does no heap allocation (#19383)

esphome/components/json/__init__.py

@@ -21,3 +21,9 @@ async def to_code(config: ConfigType) -> None:
         cg.add_library("bblanchon/ArduinoJson", "7.4.3")
     cg.add_define("USE_JSON")
     cg.add_global(json_ns.using)
+
+
+def enable_arena() -> None:
+    """Compile JsonArena and the allocator constructor of JsonBuilder; only the consumers that build
+    documents in a stack arena pay for them."""
+    cg.add_define("USE_JSON_ARENA")

esphome/components/json/json_util.cpp

@@ -46,11 +46,7 @@ JsonDocument parse_json(const uint8_t *data, size_t len) {
     ESP_LOGE(TAG, "No data to parse");
     return JsonObject();  // return unbound object
   }
-#ifdef USE_PSRAM
-  JsonDocument json_document(&global_json_allocator);
-#else
-  JsonDocument json_document;
-#endif
+  JsonDocument json_document(heap_json_allocator());
   if (json_document.overflowed()) {
     ESP_LOGE(TAG, "Could not allocate memory for JSON document!");
     return JsonObject();  // return unbound object
@@ -68,7 +64,18 @@ JsonDocument parse_json(const uint8_t *data, size_t len) {
   // NOLINTEND(clang-analyzer-cplusplus.NewDeleteLeaks,clang-analyzer-core.StackAddressEscape)
 }
 
-JsonBuilder::JsonBuilder() = default;
+JsonBuilder::JsonBuilder() : doc_(heap_json_allocator()) {}
+#ifdef USE_JSON_ARENA
+JsonBuilder::JsonBuilder(ArduinoJson::Allocator *allocator) : doc_(allocator) {}
+#endif
+
+ArduinoJson::Allocator *heap_json_allocator() {
+#ifdef USE_PSRAM
+  return &global_json_allocator;
+#else
+  return ArduinoJson::detail::DefaultAllocator::instance();
+#endif
+}
 
 size_t JsonBuilder::serialize_to(char *buf, size_t cap) {
   if (doc_.overflowed()) {

esphome/components/json/json_util.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <algorithm>
+#include <cstddef>
 #include <cstring>
 #include <string>
 #include <vector>
@@ -165,11 +167,88 @@ inline JsonDocument parse_json(const std::string &data) {
   return parse_json(reinterpret_cast<const uint8_t *>(data.c_str()), data.size());
 }
 
+/// The allocator a JsonBuilder uses by default (PSRAM first when available)
+ArduinoJson::Allocator *heap_json_allocator();
+
+#ifdef USE_JSON_ARENA
+/// Size of one ArduinoJson slot pool, the first allocation every document makes (1 KB on 32 bit targets)
+constexpr size_t JSON_POOL_BYTES = ARDUINOJSON_POOL_CAPACITY * sizeof(ArduinoJson::detail::VariantData);
+/// One pool plus room for copied string nodes; a 40 option select with linked options fits
+constexpr size_t JSON_ARENA_SIZE = JSON_POOL_BYTES + 1152;
+static_assert(sizeof(void *) != 4 || JSON_ARENA_SIZE == 2176, "the arena was sized for a 1 KB pool");
+
+/// Bump allocator over a fixed buffer for a document built and serialized in one scope. What the
+/// buffer cannot hold goes to the heap allocator; nothing is freed until the arena goes away.
+template<size_t N> class JsonArena final : public ArduinoJson::Allocator {
+ public:
+  // Takes what the buffer cannot hold
+  explicit JsonArena(ArduinoJson::Allocator *fallback = heap_json_allocator()) : fallback_(fallback) {}
+  // The document points into buf_
+  JsonArena(const JsonArena &) = delete;
+  JsonArena &operator=(const JsonArena &) = delete;
+
+  void *allocate(size_t size) override {
+    if (size > N) {
+      return this->fallback_->allocate(size);  // also keeps the rounding below from wrapping
+    }
+    size = (size + ALIGN - 1) & ~(ALIGN - 1);
+    if (size > N - this->used_) {
+      return this->fallback_->allocate(size);
+    }
+    this->last_ = this->used_;
+    this->used_ += size;
+    return this->buf_ + this->last_;
+  }
+  void deallocate(void *ptr) override {
+    if (!this->owns_(ptr)) {
+      this->fallback_->deallocate(ptr);
+    }
+  }
+  void *reallocate(void *ptr, size_t new_size) override {
+    if (!this->owns_(ptr)) {
+      return this->fallback_->reallocate(ptr, new_size);
+    }
+    const size_t off = static_cast<uint8_t *>(ptr) - this->buf_;
+    const size_t size = new_size > N ? N + ALIGN : (new_size + ALIGN - 1) & ~(ALIGN - 1);
+    const bool newest = off == this->last_;
+    if (newest && size <= N - off) {
+      this->used_ = off + size;  // the newest block grows or shrinks in place
+      return ptr;
+    }
+    // An older block's size is unknown; copying to the end of the buffer stays in bounds
+    const size_t old_size = newest ? this->used_ - off : N - off;
+    void *moved = this->fallback_->allocate(new_size);
+    if (moved == nullptr) {
+      return nullptr;  // the caller keeps ptr, so its arena space stays reserved
+    }
+    std::memcpy(moved, ptr, std::min(new_size, old_size));
+    if (newest) {
+      this->used_ = off;  // it moved to the heap, so its arena space is free again
+    }
+    return moved;
+  }
+  /// Bytes of the buffer handed out so far
+  size_t used() const { return this->used_; }
+
+ private:
+  static constexpr size_t ALIGN = alignof(std::max_align_t);
+  bool owns_(const void *ptr) const { return ptr >= this->buf_ && ptr < this->buf_ + N; }
+  ArduinoJson::Allocator *fallback_;
+  alignas(ALIGN) uint8_t buf_[N];
+  size_t used_{0};
+  size_t last_{0};
+};
+#endif  // USE_JSON_ARENA
+
 /// Builder class for creating JSON documents without lambdas
 class JsonBuilder {
  public:
   // Out of line: inlining the JsonDocument constructor duplicates it at every call site
   JsonBuilder();
+#ifdef USE_JSON_ARENA
+  // The builder must not outlive the allocator
+  explicit JsonBuilder(ArduinoJson::Allocator *allocator);
+#endif
 
   JsonObject root() {
     if (!root_created_) {
@@ -188,12 +267,7 @@ class JsonBuilder {
   SerializationBuffer<> serialize();
 
  private:
-#ifdef USE_PSRAM
-  SpiRamAllocator allocator_;
-  JsonDocument doc_{&allocator_};
-#else
   JsonDocument doc_;
-#endif
   JsonObject root_;
   bool root_created_{false};
 };

esphome/components/web_server/__init__.py

@@ -8,6 +8,7 @@ from typing import Any
 
 import esphome.codegen as cg
 from esphome.components import web_server_base
+from esphome.components.json import enable_arena
 from esphome.components.logger import request_log_listener
 from esphome.components.web_server_base import CONF_WEB_SERVER_BASE_ID
 import esphome.config_validation as cv
@@ -387,6 +388,9 @@ async def to_code(config: ConfigType) -> None:
         cg.add(paren.set_port(port))
     cg.add_define("USE_WEBSERVER")
     cg.add_define("USE_WEBSERVER_PORT", port)
+    if CORE.is_esp32:
+        # The ESP-IDF event source builds state documents in a stack arena
+        enable_arena()
     cg.add_define("USE_WEBSERVER_VERSION", version)
     if version >= 2:
         # Don't compress the index HTML as the data sizes are almost the same.

esphome/components/web_server_idf/web_server_idf.cpp

@@ -57,6 +57,12 @@ namespace esphome::web_server_idf {
 
 static const char *const TAG = "web_server_idf";
 
+// Only send_json_() may hold the JSON arena: every other frame in this file is capped below one
+// arena. Measured at -Os on GCC 14; newer toolchains stay checked on purpose, older ones skip it.
+#if defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 14 && defined(__OPTIMIZE_SIZE__)
+#pragma GCC diagnostic error "-Wstack-usage=2048"
+#endif
+
 // Chunk size for streaming request bodies; matches the Arduino AsyncWebServer buffer size.
 // Buffers of this size must live on the heap - the httpd task stack is too small.
 static constexpr size_t RECV_CHUNK_SIZE = 1460;
@@ -893,9 +899,7 @@ void AsyncEventSourceResponse::process_deferred_queue_() {
   }
   while (!deferred_queue_.empty()) {
     DeferredEvent &de = deferred_queue_.front();
-    json::JsonBuilder builder;
-    de.message_generator_(web_server_, de.source_, builder);
-    if (this->send_json_(builder)) {
+    if (this->send_json_(de.source_, de.message_generator_)) {
       if (this->close_requested_ || deferred_queue_.empty()) {
         return;
       }
@@ -1086,7 +1090,15 @@ void AsyncEventSourceResponse::loop() {
   this->entities_iterator_.try_advance(1);
 }
 
-bool AsyncEventSourceResponse::send_json_(json::JsonBuilder &builder) {
+#if defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 14 && defined(__OPTIMIZE_SIZE__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstack-usage="  // the one frame that holds the arena and the JSON buffer
+#endif
+bool AsyncEventSourceResponse::send_json_(void *source, message_generator_t *generator) {
+  // The arena lives only in this frame, so no call chain ever holds two of them
+  json::JsonArena<json::JSON_ARENA_SIZE> arena;
+  json::JsonBuilder builder(&arena);
+  generator(this->web_server_, source, builder);
   char buf[JSON_BUF_SIZE];
   const size_t len = builder.serialize_to(buf, sizeof(buf));
   if (len < sizeof(buf)) {
@@ -1138,6 +1150,9 @@ bool AsyncEventSourceResponse::send_json_(json::JsonBuilder &builder) {
   drain_tail_();
   return true;
 }
+#if defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 14 && defined(__OPTIMIZE_SIZE__)
+#pragma GCC diagnostic pop
+#endif
 
 void AsyncEventSourceResponse::tail_alloc_failed_(size_t cap) {
   // Same stall clock as a socket that stops draining, so a session cannot retry forever
@@ -1250,10 +1265,8 @@ void AsyncEventSourceResponse::deferrable_send_state(void *source, const char *e
     // trying to send first
     deq_push_back_with_dedup_(source, message_generator);
   } else {
-    json::JsonBuilder builder;
-    message_generator(web_server_, source, builder);
     // A send error closes the session and clears the queue; nothing is queued after that
-    if (!this->send_json_(builder) && !this->close_requested_) {
+    if (!this->send_json_(source, message_generator) && !this->close_requested_) {
       deq_push_back_with_dedup_(source, message_generator);
     }
   }

esphome/components/web_server_idf/web_server_idf.h

@@ -316,7 +316,7 @@ class AsyncEventSourceResponse {
   bool stash_chunk_(const char *prefix, size_t prefix_len, const char *message, size_t message_len, size_t total,
                     size_t sent);
   // Send a state event; JSON too large for the stack buffer is serialized into tail_ instead
-  bool send_json_(json::JsonBuilder &builder);
+  bool send_json_(void *source, message_generator_t *generator);
   // Warn once, and close the session once the stall timeout passes with no memory for the tail
   void tail_alloc_failed_(size_t cap);
   void request_close_();
@@ -362,6 +362,7 @@ class AsyncEventSourceResponse {
   static constexpr size_t MAX_SEND_IOV = 1 + 2 * MAX_SEND_LINES;
   // Chunk header, retry/id/event lines and the first "data: "
   static constexpr size_t PREFIX_BUF_SIZE = 128;
+
   // Stack buffer for a state event's JSON; a larger document is serialized into the tail
   static constexpr size_t JSON_BUF_SIZE = 1024;
   // Same ceiling JsonBuilder::serialize() applies (max_heap_size in json_util.cpp); a larger

esphome/core/defines.h

@@ -87,6 +87,7 @@
 #define USE_INFRARED
 #define USE_IR_RF
 #define USE_JSON
+#define USE_JSON_ARENA
 #define USE_RADIO_FREQUENCY
 #define USE_LIGHT
 #define USE_LIGHT_FLASH_TRANSITION_LENGTH

tests/components/json/__init__.py

@@ -1,3 +1,6 @@
+import functools
+
+from esphome.components.json import enable_arena
 from tests.testing_helpers import ComponentManifestOverride
 
 
@@ -7,3 +10,12 @@ def override_manifest(manifest: ComponentManifestOverride) -> None:
     # library registration to happen, otherwise json_util.cpp fails to find
     # ArduinoJson.h.
     manifest.enable_codegen()
+    # The JsonArena host test needs the arena compiled in, as a consumer would request it
+    real_to_code = manifest.to_code
+
+    @functools.wraps(real_to_code)
+    async def to_code_with_arena(config):
+        await real_to_code(config)
+        enable_arena()
+
+    manifest.to_code = to_code_with_arena

tests/components/json/test_json_arena.cpp

@@ -0,0 +1,238 @@
+#include <gtest/gtest.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+
+#include "esphome/components/json/json_util.h"
+
+using esphome::json::JsonArena;
+using esphome::json::JsonBuilder;
+
+namespace {
+
+constexpr size_t ALIGN = alignof(std::max_align_t);
+constexpr size_t round_up(size_t n) { return (n + ALIGN - 1) & ~(ALIGN - 1); }
+
+// Counts what the arena could not hold
+struct Counting final : ArduinoJson::Allocator {
+  int allocs{0};
+  void *allocate(size_t n) override {
+    this->allocs++;
+    return malloc(n);  // NOLINT
+  }
+  void deallocate(void *p) override { free(p); }  // NOLINT
+  void *reallocate(void *p, size_t n) override {
+    this->allocs++;
+    return realloc(p, n);  // NOLINT
+  }
+};
+
+// Refuses everything, so a spill or a move sees the heap as exhausted
+struct NoMemory final : ArduinoJson::Allocator {
+  void *allocate(size_t) override { return nullptr; }
+  void deallocate(void *) override {}
+  void *reallocate(void *, size_t) override { return nullptr; }
+};
+
+template<size_t N> bool inside(const JsonArena<N> &arena, const void *p) {
+  auto base = reinterpret_cast<uintptr_t>(&arena);
+  auto addr = reinterpret_cast<uintptr_t>(p);
+  return addr >= base && addr < base + sizeof(arena);
+}
+
+}  // namespace
+
+TEST(JsonArena, BumpsAlignedInsideTheBuffer) {
+  JsonArena<256> arena;
+  auto *a = static_cast<uint8_t *>(arena.allocate(10));
+  auto *b = static_cast<uint8_t *>(arena.allocate(10));
+  ASSERT_NE(a, nullptr);
+  ASSERT_NE(b, nullptr);
+  EXPECT_TRUE(inside(arena, a));
+  EXPECT_TRUE(inside(arena, b));
+  EXPECT_EQ(reinterpret_cast<uintptr_t>(a) % ALIGN, 0u);
+  EXPECT_EQ(static_cast<size_t>(b - a), round_up(10));
+  arena.deallocate(a);
+  arena.deallocate(b);
+}
+
+TEST(JsonArena, SpillsToTheHeapWhenFull) {
+  JsonArena<64> arena;
+  void *a = arena.allocate(48);
+  void *b = arena.allocate(48);
+  ASSERT_NE(a, nullptr);
+  ASSERT_NE(b, nullptr);
+  EXPECT_TRUE(inside(arena, a));
+  EXPECT_FALSE(inside(arena, b));
+  std::memset(b, 'b', 48);
+  arena.deallocate(b);  // routed to the heap; a mismatch would trip the sanitizer
+  arena.deallocate(a);
+}
+
+TEST(JsonArena, NewestBlockGrowsAndShrinksInPlace) {
+  JsonArena<256> arena;
+  void *a = arena.allocate(16);
+  std::memset(a, 'x', 16);
+  EXPECT_EQ(arena.reallocate(a, 96), a);
+  EXPECT_EQ(std::memcmp(a, "xxxxxxxxxxxxxxxx", 16), 0);
+  EXPECT_EQ(arena.reallocate(a, 8), a);
+  auto *next = static_cast<uint8_t *>(arena.allocate(8));
+  EXPECT_EQ(static_cast<size_t>(next - static_cast<uint8_t *>(a)), round_up(8));
+}
+
+TEST(JsonArena, NewestBlockMovesToTheHeapAndFreesItsSpace) {
+  JsonArena<64> arena;
+  void *a = arena.allocate(32);
+  std::memset(a, 'q', 32);
+  void *moved = arena.reallocate(a, 200);
+  ASSERT_NE(moved, nullptr);
+  EXPECT_FALSE(inside(arena, moved));
+  EXPECT_EQ(std::memcmp(moved, "qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq", 32), 0);
+  EXPECT_EQ(arena.allocate(16), a);  // the space it left is handed out again
+  arena.deallocate(moved);
+}
+
+TEST(JsonArena, OlderBlockMovesToTheHeapKeepingItsBytes) {
+  JsonArena<256> arena;
+  void *a = arena.allocate(16);
+  std::memset(a, 'a', 16);
+  auto *b = static_cast<uint8_t *>(arena.allocate(16));
+  void *moved = arena.reallocate(a, 64);
+  ASSERT_NE(moved, nullptr);
+  EXPECT_FALSE(inside(arena, moved));
+  EXPECT_EQ(std::memcmp(moved, "aaaaaaaaaaaaaaaa", 16), 0);
+  auto *next = static_cast<uint8_t *>(arena.allocate(8));
+  EXPECT_EQ(static_cast<size_t>(next - b), round_up(16));  // b's space is untouched
+  arena.deallocate(moved);
+}
+
+TEST(JsonArena, HeapBlocksReallocateOnTheHeap) {
+  JsonArena<32> arena;
+  void *a = arena.allocate(64);  // never fit
+  EXPECT_FALSE(inside(arena, a));
+  std::memset(a, 'h', 64);
+  void *grown = arena.reallocate(a, 128);
+  ASSERT_NE(grown, nullptr);
+  EXPECT_EQ(std::memcmp(grown, "hhhhhhhhhhhhhhhh", 16), 0);
+  arena.deallocate(grown);
+}
+
+TEST(JsonArena, FailedMoveKeepsTheBlockReserved) {
+  NoMemory no_memory;
+  JsonArena<64> arena(&no_memory);
+  void *a = arena.allocate(32);
+  std::memset(a, 'k', 32);
+  EXPECT_EQ(arena.reallocate(a, 200), nullptr);
+  EXPECT_EQ(std::memcmp(a, "kkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkk", 32), 0);
+  auto *b = static_cast<uint8_t *>(arena.allocate(16));  // must not hand out a's bytes again
+  ASSERT_NE(b, nullptr);
+  EXPECT_EQ(static_cast<size_t>(b - static_cast<uint8_t *>(a)), round_up(32));
+  EXPECT_EQ(arena.allocate(64), nullptr);  // nothing left and the fallback refuses
+}
+
+// NOLINTBEGIN(clang-analyzer-cplusplus.NewDeleteLeaks) false positive with ArduinoJson
+constexpr size_t ARENA = esphome::json::JSON_ARENA_SIZE;
+
+// The documents the event stream sends must fit without touching the fallback, and their copied
+// strings must land in the headroom above the pool
+TEST(JsonArena, StateDocumentsFitWithoutTouchingTheFallback) {
+  Counting counting;
+  JsonArena<ARENA> arena(&counting);
+  {
+    // A switch state event: copied id, domain and name, bool value
+    JsonBuilder builder(&arena);
+    JsonObject root = builder.root();
+    char id_buf[] = "switch/SSE Toggle";
+    char domain_buf[] = "switch";
+    char name_buf[] = "SSE Toggle";
+    root["id"] = static_cast<const char *>(id_buf);
+    root["domain"] = static_cast<const char *>(domain_buf);
+    root["name"] = static_cast<const char *>(name_buf);
+    root["icon"] = "";
+    root["entity_category"] = 0;
+    root["value"] = true;
+    root["state"] = "ON";
+    root["assumed_state"] = false;
+    char out[256];
+    EXPECT_LT(builder.serialize_to(out, sizeof(out)), sizeof(out));
+  }
+  EXPECT_EQ(counting.allocs, 0);
+  EXPECT_GT(arena.used(), esphome::json::JSON_POOL_BYTES);
+
+  Counting counting_select;
+  JsonArena<ARENA> select_arena(&counting_select);
+  {
+    // A 40 option select detail document: copied id, domain, name and value, linked options
+    JsonBuilder builder(&select_arena);
+    JsonObject root = builder.root();
+    char id_buf[] = "select/SSE Big Select";
+    char domain_buf[] = "select";
+    char name_buf[] = "SSE Big Select";
+    char value_buf[] = "option number 17 padded to twenty";
+    root["id"] = static_cast<const char *>(id_buf);
+    root["domain"] = static_cast<const char *>(domain_buf);
+    root["name"] = static_cast<const char *>(name_buf);
+    root["icon"] = "";
+    root["entity_category"] = 0;
+    root["value"] = static_cast<const char *>(value_buf);
+    root["state"] = static_cast<const char *>(value_buf);
+    JsonArray options = root["option"].to<JsonArray>();
+    char option_bufs[40][34];
+    for (int i = 0; i < 40; i++) {
+      snprintf(option_bufs[i], sizeof(option_bufs[i]), "option number %02d padded to twenty", i);
+      options.add(JsonString(option_bufs[i], true));
+    }
+    char out[2048];
+    EXPECT_LT(builder.serialize_to(out, sizeof(out)), sizeof(out));
+  }
+  EXPECT_EQ(counting_select.allocs, 0);
+  EXPECT_GT(select_arena.used(), esphome::json::JSON_POOL_BYTES);
+
+  // The same select with its 40 options copied, as the generator does before the strings are
+  // linked, does not fit: the headroom is sized for linked options and the rest spills
+  Counting counting_copied;
+  JsonArena<ARENA> copied_arena(&counting_copied);
+  {
+    JsonBuilder builder(&copied_arena);
+    JsonArray options = builder.root()["option"].to<JsonArray>();
+    char option_bufs[40][34];
+    for (int i = 0; i < 40; i++) {
+      snprintf(option_bufs[i], sizeof(option_bufs[i]), "option number %02d padded to twenty", i);
+      options.add(static_cast<const char *>(option_bufs[i]));
+    }
+    char out[2048];
+    EXPECT_LT(builder.serialize_to(out, sizeof(out)), sizeof(out));
+  }
+  EXPECT_GT(counting_copied.allocs, 0);
+  EXPECT_GT(copied_arena.used(), ARENA - 64);  // the arena filled up before the spill began
+}
+
+TEST(JsonArena, DocumentMatchesTheHeapAllocator) {
+  // 700 integers need six pools, which also grows ArduinoJson's pool list past its preallocated four
+  auto build = [](JsonBuilder &builder) {
+    JsonArray arr = builder.root()["a"].to<JsonArray>();
+    for (int i = 0; i < 700; i++) {
+      arr.add(i);
+    }
+    JsonArray strings = builder.root()["s"].to<JsonArray>();
+    char buf[32];
+    for (int i = 0; i < 60; i++) {
+      snprintf(buf, sizeof(buf), "string number %04d padded", i);
+      strings.add(buf);
+    }
+  };
+  JsonArena<ARENA> arena;
+  JsonBuilder with_arena(&arena);
+  build(with_arena);
+  JsonBuilder with_heap;
+  build(with_heap);
+  std::string a = with_arena.serialize();
+  std::string b = with_heap.serialize();
+  EXPECT_GT(a.size(), 4000u);
+  EXPECT_EQ(a, b);
+}
+// NOLINTEND(clang-analyzer-cplusplus.NewDeleteLeaks)