[esp32] Keep the NVS cache in PSRAM when PSRAM is guaranteed (#19378)

esphome/components/esp32/__init__.py

@@ -11,7 +11,10 @@ from typing import Any
 
 from esphome import yaml_util
 import esphome.codegen as cg
-from esphome.components.const import CONF_ENABLE_OTA_DOWNGRADE_PROTECTION
+from esphome.components.const import (
+    CONF_ENABLE_OTA_DOWNGRADE_PROTECTION,
+    CONF_IGNORE_NOT_FOUND,
+)
 from esphome.config_helpers import filter_source_files_from_defines
 import esphome.config_validation as cv
 from esphome.const import (
@@ -111,6 +114,7 @@ CONF_ENGINEERING_SAMPLE = "engineering_sample"
 CONF_INCLUDE_BUILTIN_IDF_COMPONENTS = "include_builtin_idf_components"
 CONF_ENABLE_LWIP_ASSERT = "enable_lwip_assert"
 CONF_EXECUTE_FROM_PSRAM = "execute_from_psram"
+CONF_NVS_CACHE_IN_PSRAM = "nvs_cache_in_psram"
 CONF_FLASH_CHIP = "flash_chip"
 CONF_KEY_ID = "key_id"
 CONF_MINIMUM_CHIP_REVISION = "minimum_chip_revision"
@@ -1595,6 +1599,29 @@ def final_validate(config) -> None:
                     path=[CONF_FRAMEWORK, CONF_ADVANCED, CONF_EXECUTE_FROM_PSRAM],
                 )
             )
+    if advanced.get(CONF_NVS_CACHE_IN_PSRAM):
+        psram_conf = full_config.get(PSRAM_DOMAIN)
+        if (
+            psram_conf is None
+            or psram_conf[CONF_DISABLED]
+            or psram_conf[CONF_IGNORE_NOT_FOUND]
+        ):
+            errs.append(
+                cv.Invalid(
+                    f"'{CONF_NVS_CACHE_IN_PSRAM}' requires PSRAM with 'ignore_not_found: false'",
+                    path=[CONF_FRAMEWORK, CONF_ADVANCED, CONF_NVS_CACHE_IN_PSRAM],
+                )
+            )
+        if (
+            advanced.get(CONF_NVS_ENCRYPTION) is not None
+            or conf_fw[CONF_SDKCONFIG_OPTIONS].get("CONFIG_NVS_ENCRYPTION") == "y"
+        ):
+            errs.append(
+                cv.Invalid(
+                    f"'{CONF_NVS_CACHE_IN_PSRAM}' cannot be used with NVS encryption; the keys must stay in internal RAM",
+                    path=[CONF_FRAMEWORK, CONF_ADVANCED, CONF_NVS_CACHE_IN_PSRAM],
+                )
+            )
 
     final_validate_pins(full_config)
 
@@ -2035,6 +2062,7 @@ FRAMEWORK_SCHEMA = cv.Schema(
                 cv.Optional(CONF_RINGBUF_IN_IRAM, default=False): cv.boolean,
                 cv.Optional(CONF_HEAP_IN_IRAM, default=False): cv.boolean,
                 cv.Optional(CONF_EXECUTE_FROM_PSRAM, default=False): cv.boolean,
+                cv.Optional(CONF_NVS_CACHE_IN_PSRAM): cv.boolean,
                 cv.Optional(CONF_LOOP_TASK_STACK_SIZE, default=8192): cv.int_range(
                     min=8192, max=32768
                 ),
@@ -2361,6 +2389,20 @@ async def _set_libc_picolibc_newlib_compat() -> None:
     )
 
 
+@coroutine_with_priority(CoroPriority.FINAL)
+async def _apply_nvs_cache_in_psram(explicit: bool) -> None:
+    """Keep the NVS cache in PSRAM unless NVS encryption is on, however it was enabled."""
+    # The encrypted partition object holds the derived keys, which must stay in internal RAM
+    if is_idf_sdkconfig_option_enabled("CONFIG_NVS_ENCRYPTION"):
+        if explicit:
+            _LOGGER.warning(
+                "%s ignored: NVS encryption keeps the NVS cache in internal RAM",
+                CONF_NVS_CACHE_IN_PSRAM,
+            )
+        return
+    set_idf_sdkconfig_default("CONFIG_NVS_ALLOCATE_CACHE_IN_SPIRAM", True)
+
+
 @coroutine_with_priority(CoroPriority.FINAL)
 async def _reconcile_certificate_bundle_sdkconfig() -> None:
     """Enable the mbedTLS certificate bundle only when something asked for it.
@@ -2913,6 +2955,17 @@ async def to_code(config):
     if advanced[CONF_EXECUTE_FROM_PSRAM]:
         add_idf_sdkconfig_option("CONFIG_SPIRAM_XIP_FROM_PSRAM", True)
 
+    # Imported here as psram imports this module
+    from esphome.components.psram import is_guaranteed as psram_is_guaranteed
+
+    # Frees internal heap (the cache scales with the NVS partition) but slows NVS, so only
+    # where PSRAM is known to be fitted. Decided at FINAL so every way of enabling NVS
+    # encryption has been seen and a user's sdkconfig_options value wins.
+    # Unset means on; only an explicit true is worth a warning when it has to be dropped.
+    requested = advanced.get(CONF_NVS_CACHE_IN_PSRAM)
+    if requested is not False and psram_is_guaranteed():
+        CORE.add_job(_apply_nvs_cache_in_psram, requested is True)
+
     # Apply LWIP core locking for better socket performance
     # This is already enabled by default in Arduino framework, where it provides
     # significant performance benefits. Our benchmarks show socket operations are

tests/component_tests/esp32/config/nvs_cache_psram_disabled.yaml

@@ -0,0 +1,13 @@
+esphome:
+  name: test
+
+esp32:
+  variant: esp32s3
+  framework:
+    type: esp-idf
+    advanced:
+      nvs_cache_in_psram: false
+
+psram:
+  mode: octal
+  ignore_not_found: false

tests/component_tests/esp32/config/nvs_cache_psram_encrypted.yaml

@@ -0,0 +1,14 @@
+esphome:
+  name: test
+
+esp32:
+  variant: esp32s3
+  framework:
+    type: esp-idf
+    advanced:
+      nvs_encryption:
+        key_id: 0
+
+psram:
+  mode: octal
+  ignore_not_found: false

tests/component_tests/esp32/config/nvs_cache_psram_encrypted_sdkconfig.yaml

@@ -0,0 +1,13 @@
+esphome:
+  name: test
+
+esp32:
+  variant: esp32s3
+  framework:
+    type: esp-idf
+    sdkconfig_options:
+      CONFIG_NVS_ENCRYPTION: "y"
+
+psram:
+  mode: octal
+  ignore_not_found: false

tests/component_tests/esp32/config/nvs_cache_psram_explicit.yaml

@@ -0,0 +1,13 @@
+esphome:
+  name: test
+
+esp32:
+  variant: esp32s3
+  framework:
+    type: esp-idf
+    advanced:
+      nvs_cache_in_psram: true
+
+psram:
+  mode: octal
+  ignore_not_found: false

tests/component_tests/esp32/config/nvs_cache_psram_guaranteed.yaml

@@ -0,0 +1,11 @@
+esphome:
+  name: test
+
+esp32:
+  variant: esp32s3
+  framework:
+    type: esp-idf
+
+psram:
+  mode: octal
+  ignore_not_found: false

tests/component_tests/esp32/config/nvs_cache_psram_not_guaranteed.yaml

@@ -0,0 +1,10 @@
+esphome:
+  name: test
+
+esp32:
+  variant: esp32s3
+  framework:
+    type: esp-idf
+
+psram:
+  mode: octal

tests/component_tests/esp32/config/nvs_cache_psram_user_off.yaml

@@ -0,0 +1,13 @@
+esphome:
+  name: test
+
+esp32:
+  variant: esp32s3
+  framework:
+    type: esp-idf
+    sdkconfig_options:
+      CONFIG_NVS_ALLOCATE_CACHE_IN_SPIRAM: "n"
+
+psram:
+  mode: octal
+  ignore_not_found: false

tests/component_tests/esp32/test_esp32.py

@@ -4,6 +4,7 @@ Test ESP32 configuration
 
 import asyncio
 from collections.abc import Callable
+import logging
 from pathlib import Path
 from typing import Any
 
@@ -1575,3 +1576,137 @@ def test_vasprintf_stub_only_on_rom_vsnprintf_variants(
     assert (CORE.build_flags >= _VASPRINTF_STUB_FLAGS) is expected
     defines = {define.name for define in CORE.defines}
     assert ("USE_ESP32_VASPRINTF_STUB" in defines) is expected
+
+
+@pytest.mark.parametrize(
+    ("fixture", "expected"),
+    [
+        ("nvs_cache_psram_guaranteed.yaml", True),
+        ("nvs_cache_psram_explicit.yaml", True),
+        ("nvs_cache_psram_not_guaranteed.yaml", None),
+        ("nvs_cache_psram_disabled.yaml", None),
+        # the encryption keys must stay in internal RAM, whichever way encryption is enabled
+        ("nvs_cache_psram_encrypted.yaml", None),
+        ("nvs_cache_psram_encrypted_sdkconfig.yaml", None),
+    ],
+)
+def test_nvs_cache_in_psram_sdkconfig(
+    generate_main: Callable[[str | Path], str],
+    component_config_path: Callable[[str], Path],
+    fixture: str,
+    expected: bool | None,
+) -> None:
+    """The NVS cache moves to PSRAM only with guaranteed PSRAM, the option not off and no NVS encryption."""
+    generate_main(component_config_path(fixture))
+    sdkconfig = CORE.data[KEY_ESP32][KEY_SDKCONFIG_OPTIONS]
+    assert sdkconfig.get("CONFIG_NVS_ALLOCATE_CACHE_IN_SPIRAM") is expected
+
+
+def test_nvs_cache_in_psram_user_sdkconfig_wins(
+    generate_main: Callable[[str | Path], str],
+    component_config_path: Callable[[str], Path],
+) -> None:
+    """A raw sdkconfig_options value for the NVS cache option is left alone."""
+    generate_main(component_config_path("nvs_cache_psram_user_off.yaml"))
+    sdkconfig = CORE.data[KEY_ESP32][KEY_SDKCONFIG_OPTIONS]
+    assert sdkconfig["CONFIG_NVS_ALLOCATE_CACHE_IN_SPIRAM"] == RawSdkconfigValue("n")
+
+
+@pytest.mark.parametrize(
+    ("full_config", "error_match"),
+    [
+        pytest.param(
+            {CONF_ESPHOME: {}, "psram": {"disabled": False, "ignore_not_found": True}},
+            r"'nvs_cache_in_psram' requires PSRAM with 'ignore_not_found: false'",
+            id="nvs_cache_in_psram_needs_guaranteed_psram",
+        ),
+        pytest.param(
+            {CONF_ESPHOME: {}},
+            r"'nvs_cache_in_psram' requires PSRAM with 'ignore_not_found: false'",
+            id="nvs_cache_in_psram_needs_psram",
+        ),
+    ],
+)
+def test_nvs_cache_in_psram_explicit_true_errors(
+    full_config: dict, error_match: str, set_core_config: SetCoreConfigCallable
+) -> None:
+    """An explicit nvs_cache_in_psram: true that cannot apply is a config error, not a silent no-op."""
+    set_core_config(PlatformFramework.ESP32_IDF, full_config=full_config)
+    from esphome.components.esp32 import CONFIG_SCHEMA, FINAL_VALIDATE_SCHEMA
+
+    config = {
+        "variant": "esp32s3",
+        "framework": {"type": "esp-idf", "advanced": {"nvs_cache_in_psram": True}},
+    }
+    with pytest.raises(cv.Invalid, match=error_match):
+        FINAL_VALIDATE_SCHEMA(CONFIG_SCHEMA(config))
+
+
+def test_nvs_cache_in_psram_explicit_true_rejects_encryption(
+    set_core_config: SetCoreConfigCallable,
+) -> None:
+    set_core_config(
+        PlatformFramework.ESP32_IDF,
+        full_config={
+            CONF_ESPHOME: {},
+            "psram": {"disabled": False, "ignore_not_found": False},
+        },
+    )
+    from esphome.components.esp32 import CONFIG_SCHEMA, FINAL_VALIDATE_SCHEMA
+
+    config = {
+        "variant": "esp32s3",
+        "framework": {
+            "type": "esp-idf",
+            "advanced": {"nvs_cache_in_psram": True, "nvs_encryption": {"key_id": 0}},
+        },
+    }
+    with pytest.raises(cv.Invalid, match="cannot be used with NVS encryption"):
+        FINAL_VALIDATE_SCHEMA(CONFIG_SCHEMA(config))
+
+
+def test_nvs_cache_in_psram_default_with_encryption_is_quiet(
+    generate_main: Callable[[str | Path], str],
+    component_config_path: Callable[[str], Path],
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Encryption on a board that never mentioned the option must not warn about it."""
+    with caplog.at_level(logging.WARNING):
+        generate_main(component_config_path("nvs_cache_psram_encrypted.yaml"))
+    assert "nvs_cache_in_psram" not in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_nvs_cache_in_psram_explicit_request_warns_when_encrypted(
+    set_core_config: SetCoreConfigCallable, caplog: pytest.LogCaptureFixture
+) -> None:
+    """An explicit request dropped for NVS encryption enabled elsewhere logs a warning."""
+    set_core_config(
+        PlatformFramework.ESP32_IDF, platform_data={KEY_SDKCONFIG_OPTIONS: {}}
+    )
+    from esphome.components.esp32 import (
+        _apply_nvs_cache_in_psram,
+        add_idf_sdkconfig_option,
+    )
+
+    add_idf_sdkconfig_option("CONFIG_NVS_ENCRYPTION", True)
+    with caplog.at_level(logging.WARNING):
+        await _apply_nvs_cache_in_psram(True)
+    assert "nvs_cache_in_psram ignored" in caplog.text
+    assert (
+        "CONFIG_NVS_ALLOCATE_CACHE_IN_SPIRAM"
+        not in CORE.data[KEY_ESP32][KEY_SDKCONFIG_OPTIONS]
+    )
+
+
+def test_nvs_cache_in_psram_explicit_true_on_valid_board_is_quiet(
+    generate_main: Callable[[str | Path], str],
+    component_config_path: Callable[[str], Path],
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """An explicit true that applies sets the option and warns about nothing."""
+    with caplog.at_level(logging.WARNING):
+        generate_main(component_config_path("nvs_cache_psram_explicit.yaml"))
+    sdkconfig = CORE.data[KEY_ESP32][KEY_SDKCONFIG_OPTIONS]
+    assert sdkconfig.get("CONFIG_NVS_ALLOCATE_CACHE_IN_SPIRAM") is True
+    assert "nvs_cache_in_psram" not in caplog.text