zephyr-cp: fix flash scratch page bugs; use 4 KB NOR flash layout pages; fix RP2 layout

ports/zephyr-cp/internal-modules/circuitpython-kconfig/Kconfig

@@ -64,6 +64,42 @@ config ZERO_LATENCY_IRQS
 config BT_CTLR_ASSERT_OVERHEAD_START
 	default n if SOC_FAMILY_NORDIC_NRF
 
+# ===== Flash filesystem scratch page =====
+# supervisor/flash.c caches the erase page being written in RAM. With a
+# scratch page, it allocates the cache only while writing, and stages writes
+# in the last erase page of the filesystem when it can't. Without one, it
+# allocates the cache once at startup and keeps it, and the filesystem gets
+# the whole partition.
+#
+# ports/raspberrypi caches one 4 KB sector in a static buffer and reserves no
+# scratch page, so do the same on RP2040 and RP2350: the filesystem sizes
+# then match, and either port can use a filesystem formatted by the other.
+config CIRCUITPY_FLASH_SCRATCH_PAGE
+	bool "Stage flash writes in a scratch page when the RAM cache can't be allocated"
+	default n if SOC_FAMILY_RPI_PICO
+	default y
+
+# ===== External NOR flash =====
+# supervisor/flash.c sets aside the last erase page of the filesystem as a
+# scratch page, and caches a whole erase page in RAM while writing. The NOR
+# flash drivers' default layout page is the 64 KB erase block, which loses
+# 64 KB of the filesystem and needs a 64 KB RAM cache. Use the 4 KB sector
+# instead; the drivers erase 4 KB sectors too. Changing one block then
+# rewrites 4 KB instead of 64 KB, which is faster for small writes such as
+# FAT updates. Filesystem partitions also start and end on 4 KB boundaries,
+# but not always on 64 KB ones.
+#
+# On nRF QSPI boards, the filesystem size then matches ports/nordic, so
+# either port can use a filesystem formatted by the other.
+config NORDIC_QSPI_NOR_FLASH_LAYOUT_PAGE_SIZE
+	default 4096 if NORDIC_QSPI_NOR
+
+config SPI_NOR_FLASH_LAYOUT_PAGE_SIZE
+	default 4096 if SPI_NOR
+
+config FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE
+	default 4096 if FLASH_MSPI_NOR
+
 # ===== Bluetooth defaults =====
 
 # Use a variable for the chosen name so the comma isn't parsed as an argument separator

ports/zephyr-cp/supervisor/flash.c

@@ -31,6 +31,28 @@
 #elif FIXED_PARTITION_EXISTS(circuitpy_partition)
 #define CIRCUITPY_PARTITION circuitpy_partition
 #endif
+
+// Check that the filesystem partition starts and ends on a layout page
+// boundary: flash.c addresses erase pages relative to the start of the
+// partition. The NOR flash drivers set their layout page size in Kconfig, so
+// they can be checked at build time. (Other drivers report it only at run
+// time.)
+#if FIXED_PARTITION_EXISTS(CIRCUITPY_PARTITION)
+#define CIRCUITPY_PARTITION_NODE DT_NODELABEL(CIRCUITPY_PARTITION)
+#define CIRCUITPY_PARTITION_MTD DT_MTD_FROM_FIXED_PARTITION(CIRCUITPY_PARTITION_NODE)
+#if defined(CONFIG_NORDIC_QSPI_NOR) && DT_NODE_HAS_COMPAT(CIRCUITPY_PARTITION_MTD, nordic_qspi_nor)
+#define CIRCUITPY_LAYOUT_PAGE_SIZE CONFIG_NORDIC_QSPI_NOR_FLASH_LAYOUT_PAGE_SIZE
+#elif defined(CONFIG_SPI_NOR) && DT_NODE_HAS_COMPAT(CIRCUITPY_PARTITION_MTD, jedec_spi_nor)
+#define CIRCUITPY_LAYOUT_PAGE_SIZE CONFIG_SPI_NOR_FLASH_LAYOUT_PAGE_SIZE
+#elif defined(CONFIG_FLASH_MSPI_NOR) && DT_NODE_HAS_COMPAT(CIRCUITPY_PARTITION_MTD, jedec_nor)
+#define CIRCUITPY_LAYOUT_PAGE_SIZE CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE
+#endif
+#ifdef CIRCUITPY_LAYOUT_PAGE_SIZE
+BUILD_ASSERT(DT_REG_ADDR(CIRCUITPY_PARTITION_NODE) % CIRCUITPY_LAYOUT_PAGE_SIZE == 0 &&
+    DT_REG_SIZE(CIRCUITPY_PARTITION_NODE) % CIRCUITPY_LAYOUT_PAGE_SIZE == 0,
+    "filesystem partition must start and end on a flash layout page boundary");
+#endif
+#endif
 static const struct flash_area *filesystem_area = NULL;
 
 #if !FIXED_PARTITION_EXISTS(CIRCUITPY_PARTITION)
@@ -58,6 +80,12 @@ static uint32_t _current_page_address;
 
 static uint32_t _scratch_page_address;
 
+// Size of the filesystem in bytes: every whole page, less the scratch page if
+// there is one.
+static uint32_t _filesystem_size;
+
+static bool allocate_ram_cache(void);
+
 // Per-block flags packed into a uint8_t array (one byte per block).
 #define ROW_LOADED  0x01  // Block data is present in the cache.
 #define ROW_DIRTY   0x02  // Block has been modified since last flush.
@@ -212,8 +240,6 @@ void supervisor_flash_init(void) {
     const struct flash_parameters *fp = flash_get_parameters(d);
     _erase_value = fp->erase_value;
     flash_get_page_info_by_offs(d, filesystem_area->fa_off, &first_info);
-    struct flash_pages_info last_info;
-    flash_get_page_info_by_offs(d, filesystem_area->fa_off + filesystem_area->fa_size - FILESYSTEM_BLOCK_SIZE, &last_info);
     _page_size = first_info.size;
     if (_page_size < FILESYSTEM_BLOCK_SIZE) {
         _page_size = FILESYSTEM_BLOCK_SIZE;
@@ -223,6 +249,17 @@ void supervisor_flash_init(void) {
     printk("  erase page size %d\n", _page_size);
     printk("  blocks per page %d\n", _blocks_per_page);
 
+    // Page addresses are relative to the start of the area. If the area
+    // doesn't end on a page boundary, its last partial page can't be erased
+    // without going past the end, so it is left unused.
+    size_t page_count = filesystem_area->fa_size / _page_size;
+    size_t min_page_count = IS_ENABLED(CONFIG_CIRCUITPY_FLASH_SCRATCH_PAGE) ? 2 : 1;
+    if (page_count < min_page_count) {
+        printk("Filesystem area too small (%zu pages)\n", page_count);
+        filesystem_area = NULL;
+        return;
+    }
+
     _row_flags = port_malloc(_blocks_per_page, false);
     if (_row_flags == NULL) {
         printk("Unable to allocate row flags (%d bytes)\n", _blocks_per_page);
@@ -231,9 +268,22 @@ void supervisor_flash_init(void) {
     }
     memset(_row_flags, 0, _blocks_per_page);
 
-    // The last page is the scratch sector.
-    _scratch_page_address = last_info.start_offset - filesystem_area->fa_off;
     _current_page_address = NO_PAGE_LOADED;
+
+    #if CONFIG_CIRCUITPY_FLASH_SCRATCH_PAGE
+    // The last whole page is the scratch page.
+    _scratch_page_address = (page_count - 1) * _page_size;
+    _filesystem_size = _scratch_page_address;
+    #else
+    // No scratch page: allocate the RAM cache now and keep it, so writes
+    // never need one.
+    if (!allocate_ram_cache()) {
+        printk("Unable to allocate the filesystem ram cache\n");
+        filesystem_area = NULL;
+        return;
+    }
+    _filesystem_size = page_count * _page_size;
+    #endif
 }
 
 uint32_t supervisor_flash_get_block_size(void) {
@@ -244,7 +294,7 @@ uint32_t supervisor_flash_get_block_count(void) {
     if (filesystem_area == NULL) {
         return 0;
     }
-    return _scratch_page_address / FILESYSTEM_BLOCK_SIZE;
+    return _filesystem_size / FILESYSTEM_BLOCK_SIZE;
 }
 
 
@@ -355,6 +405,10 @@ static bool flush_scratch_flash(void) {
         copy_block(_scratch_page_address + i * FILESYSTEM_BLOCK_SIZE,
             _current_page_address + i * FILESYSTEM_BLOCK_SIZE);
     }
+    // Every scratch block is now written, so the next write must erase the
+    // scratch page first. Unload the page so that it does.
+    _current_page_address = NO_PAGE_LOADED;
+    _clear_row_flags();
     return true;
 }
 
@@ -468,6 +522,10 @@ static bool flush_ram_cache(bool keep_cache) {
 // Delegates to the correct flash flush method depending on the existing cache.
 // TODO Don't blink the status indicator if we don't actually do any writing (hard to tell right now).
 static void _flush_keep_cache(bool keep_cache) {
+    #if !CONFIG_CIRCUITPY_FLASH_SCRATCH_PAGE
+    // Without a scratch page, the RAM cache is never released.
+    keep_cache = true;
+    #endif
     #ifdef MICROPY_HW_LED_MSC
     port_pin_set_output_level(MICROPY_HW_LED_MSC, true);
     #endif
@@ -541,8 +599,14 @@ static bool _flash_write_block(const uint8_t *data, uint32_t block) {
     // Mask out the lower bits that designate the address within the sector.
     uint32_t page_address = address & _page_mask;
     size_t block_index = (address / FILESYSTEM_BLOCK_SIZE) % _blocks_per_page;
-    // Flush the cache if we're moving onto a different page.
-    if (_current_page_address != page_address) {
+    // Flush the cache if we're moving onto a different page, or if we're
+    // writing the same block again while caching in the scratch page: flash
+    // can't be rewritten without an erase, so the second write must start a
+    // fresh scratch page.
+    bool scratch_rewrite = flash_cache_table == NULL &&
+        _current_page_address == page_address &&
+        (_row_flags[block_index] & ROW_DIRTY);
+    if (_current_page_address != page_address || scratch_rewrite) {
         // Check to see if we'd write to an erased block and aren't writing to
         // our cache. In that case we can write directly. Do not take this fast
         // path when the target block is part of the currently cached page: the

ports/zephyr-cp/tests/conftest.py

@@ -213,9 +213,9 @@ def _lfs_geometry(board, erase_page_size):
     geometry = LITTLEFS_FILESYSTEM_BOARDS[board]
     block_size = geometry["block_size"]
     page_size = max(erase_page_size, block_size)
-    # The last erase page of the partition is the supervisor flash cache's
-    # scratch page, so littlefs only gets the blocks before it.
-    last_page_start = ((geometry["partition_size"] - block_size) // page_size) * page_size
+    # The last whole erase page of the partition is the supervisor flash
+    # cache's scratch page, so littlefs only gets the blocks before it.
+    last_page_start = (geometry["partition_size"] // page_size - 1) * page_size
     return block_size, last_page_start // block_size
 
 

ports/zephyr-cp/tests/test_flash.py

@@ -192,3 +192,92 @@ def test_flash_256k_overwrite(circuitpython, board):
         circuitpython.flash_file, "overwrite.txt", board, erase_page_size=262144
     )
     assert content == "second version"
+
+
+SCRATCH_PAGE_CODE = """\
+import gc
+import storage
+storage.remount("/", readonly=False)
+
+# Open the files and build their contents first: nothing that allocates can
+# run once the heap is full.
+names = ("/scratch0.txt", "/scratch1.txt", "/scratch2.txt", "/scratch3.txt")
+chunks = (b"A" * 300, b"B" * 300, b"C" * 300, b"D" * 300)
+files = [open(name, "wb") for name in names]
+
+# Fill the heap so flash.c can't allocate its RAM cache and stages writes in
+# the scratch page instead.
+filler = []
+for size in (65536, 4096, 512, 64):
+    while True:
+        try:
+            filler.append(bytearray(size))
+        except MemoryError:
+            break
+
+# Each flush() rewrites the directory block, and appends rewrite the FAT, so
+# the same blocks are written again while their page is in the scratch page.
+for _ in range(8):
+    for i in range(4):
+        files[i].write(chunks[i])
+        files[i].flush()
+for f in files:
+    f.close()
+
+filler = None
+gc.collect()
+ok = True
+for i in range(4):
+    with open(names[i], "rb") as f:
+        if f.read() != chunks[i] * 8:
+            ok = False
+            print("MISMATCH", names[i])
+if ok:
+    print("scratch_page_ok")
+print("done")
+"""
+
+
+def check_scratch_page_run(circuitpython, board, erase_page_size):
+    """Check a SCRATCH_PAGE_CODE run, on the console and in the flash image."""
+    circuitpython.wait_until_done()
+    output = circuitpython.serial.all_output
+    debug_output = circuitpython.debug_serial.all_output
+    assert "Unable to allocate ram cache" in debug_output, "scratch page path not taken"
+    assert "Erase of" not in debug_output
+    assert "flash write failed" not in debug_output
+    assert "scratch_page_ok" in output
+
+    for i, letter in enumerate("ABCD"):
+        content = read_file_from_flash(
+            circuitpython.flash_file, f"scratch{i}.txt", board, erase_page_size=erase_page_size
+        )
+        assert content == letter * 300 * 8, f"scratch{i}.txt mismatch"
+
+
+@pytest.mark.parametrize("board", ["native_native_sim", "native_native_sim_asan"], indirect=True)
+@pytest.mark.circuitpy_drive({"code.py": SCRATCH_PAGE_CODE})
+@pytest.mark.flash_config(erase_block_size=16384)
+def test_flash_scratch_page_rewrites(circuitpython, board):
+    """Test rewriting blocks while writes are staged in the scratch page.
+
+    The flash simulator rejects writes to flash that isn't erased, so a block
+    programmed twice without an erase fails instead of silently corrupting.
+    16KB erase blocks make the RAM cache too big to allocate from a full heap,
+    and evenly divide the 0x1fc000 filesystem partition. (littlefs doesn't
+    stage writes in the scratch page.)
+    """
+    check_scratch_page_run(circuitpython, board, erase_page_size=16384)
+
+
+@pytest.mark.parametrize("board", ["native_native_sim", "native_native_sim_asan"], indirect=True)
+@pytest.mark.circuitpy_drive({"code.py": SCRATCH_PAGE_CODE})
+@pytest.mark.flash_config(erase_block_size=65536)
+def test_flash_scratch_page_partial_last_page(circuitpython, board):
+    """Test the scratch page when the partition doesn't end on a page boundary.
+
+    The 0x1fc000 filesystem partition isn't a multiple of 64KB, so its last
+    page is partial and can't be erased. The scratch page must be the last
+    whole page instead.
+    """
+    check_scratch_page_run(circuitpython, board, erase_page_size=65536)