pw_kvs: Add method to force updating of entries to new format

- Add method that forces update of any entries not on the primary
  (first) format. This update rewrites the entry and all its copies with
  the new format (magic and checksum).
- Add updating entries as a task done during FullMaintenance.
- Fix bug in CopyEntryToSector verify after write. The verify was
  looking at the source address, not the destination address.

Change-Id: I923ace9a881c264a60e2bb9fc89d9a04a66b6e1b
diff --git a/pw_kvs/entry.cc b/pw_kvs/entry.cc
index c45fab3..46f9438 100644
--- a/pw_kvs/entry.cc
+++ b/pw_kvs/entry.cc
@@ -112,7 +112,7 @@
 }
 
 StatusWithSize Entry::Copy(Address new_address) const {
-  PW_LOG_DEBUG("Copying entry from 0x%x to 0x%x as ID %" PRIu32,
+  PW_LOG_DEBUG("Copying entry from %u to %u as ID %" PRIu32,
                unsigned(address()),
                unsigned(new_address),
                transaction_id());
diff --git a/pw_kvs/entry_test.cc b/pw_kvs/entry_test.cc
index 1a5baff..93567d5 100644
--- a/pw_kvs/entry_test.cc
+++ b/pw_kvs/entry_test.cc
@@ -344,14 +344,14 @@
   EXPECT_FALSE(entry_.deleted());
 }
 
-TEST_F(EntryInFlash, Update_ReadError_NoChecksumIsOkay) {
+TEST_F(EntryInFlash, Update_ReadError_WithChecksumIsError) {
   flash_.InjectReadError(FlashError::Unconditional(Status::ABORTED));
 
   EXPECT_EQ(Status::ABORTED,
             entry_.Update(kFormatWithChecksum, kTransactionId1 + 1));
 }
 
-TEST_F(EntryInFlash, Update_ReadError_WithChecksumIsError) {
+TEST_F(EntryInFlash, Update_ReadError_NoChecksumIsOkay) {
   flash_.InjectReadError(FlashError::Unconditional(Status::ABORTED));
 
   EXPECT_EQ(Status::OK, entry_.Update(kNoChecksum, kTransactionId1 + 1));
@@ -381,37 +381,43 @@
   return value;
 }
 
-TEST_F(EntryInFlash, UpdateAndCopy_DifferentChecksum_UpdatesToNewFormat) {
-  static class Sum final : public ChecksumAlgorithm {
-   public:
-    Sum() : ChecksumAlgorithm(as_bytes(span(&state_, 1))), state_(0) {}
+class ChecksumSummation final : public ChecksumAlgorithm {
+ public:
+  ChecksumSummation() : ChecksumAlgorithm(as_bytes(span(&sum_, 1))), sum_(0) {}
 
-    void Update(span<const byte> data) override {
-      state_ = ByteSum(data, state_);
+  void Reset() override { sum_ = 0; }
+
+  void Update(span<const std::byte> data) override {
+    for (const std::byte data_byte : data) {
+      sum_ += unsigned(data_byte);
     }
+  }
 
-    void Reset() override { state_ = 0; }
+ private:
+  uint32_t sum_;
+} sum_checksum;
 
-   private:
-    uint32_t state_;
-  } sum_checksum;
+constexpr uint32_t kMagicWithSum = 0x12345678;
+constexpr EntryFormat kFormatWithSum{kMagicWithSum, &sum_checksum};
+constexpr internal::EntryFormats kFormatsWithSum(kFormatWithSum);
 
-  constexpr EntryFormat sum_format{.magic = 0x12345678,
-                                   .checksum = &sum_checksum};
+TEST_F(EntryInFlash, UpdateAndCopy_DifferentChecksum_UpdatesToNewFormat) {
+  constexpr EntryFormat kFormatWithSum{.magic = 0x12345678,
+                                       .checksum = &sum_checksum};
 
-  ASSERT_EQ(Status::OK, entry_.Update(sum_format, kTransactionId1 + 9));
+  ASSERT_EQ(Status::OK, entry_.Update(kFormatWithSum, kTransactionId1 + 9));
 
   auto result = entry_.Copy(kEntry1.size());
   ASSERT_EQ(Status::OK, result.status());
   EXPECT_EQ(kEntry1.size(), result.size());
 
   constexpr uint32_t checksum =
-      ByteSum(AsBytes(sum_format.magic)) + 0 /* checksum */ +
+      ByteSum(AsBytes(kFormatWithSum.magic)) + 0 /* checksum */ +
       0 /* alignment */ + kKey1.size() + kValue1.size() +
       ByteSum(AsBytes(kTransactionId1 + 9)) + ByteSum(kKey1) + ByteSum(kValue1);
 
   constexpr auto kNewHeader1 =
-      AsBytes(sum_format.magic,          // magic
+      AsBytes(kFormatWithSum.magic,      // magic
               checksum,                  // checksum (byte sum)
               uint8_t(0),                // alignment (changed to 16 B from 32)
               uint8_t(kKey1.size()),     // key length
@@ -449,6 +455,66 @@
                         kNewEntry1.size()));
 }
 
+TEST_F(EntryInFlash, UpdateAndCopy_DifferentFormat_UpdateAndReadBack) {
+  ASSERT_EQ(Status::OK, entry_.Update(kFormatWithSum, kTransactionId1 + 6));
+
+  FlashPartition::Address new_address = entry_.size();
+
+  StatusWithSize copy_result = entry_.Copy(new_address);
+  ASSERT_EQ(Status::OK, copy_result.status());
+  ASSERT_EQ(kEntry1.size(), copy_result.size());
+
+  Entry entry;
+  ASSERT_EQ(Status::OK,
+            Entry::Read(partition_, new_address, kFormatsWithSum, &entry));
+
+  EXPECT_EQ(Status::OK, entry.VerifyChecksumInFlash());
+  EXPECT_EQ(kFormatWithSum.magic, entry.magic());
+  EXPECT_EQ(new_address, entry.address());
+  EXPECT_EQ(kTransactionId1 + 6, entry.transaction_id());
+  EXPECT_FALSE(entry.deleted());
+}
+
+TEST_F(EntryInFlash,
+       UpdateAndCopy_DifferentFormat_UpdateFormatAndCopyMultiple) {
+  ASSERT_EQ(Status::OK, entry_.Update(kFormatWithSum, kTransactionId1 + 6));
+
+  FlashPartition::Address new_address = entry_.size();
+
+  for (int i = 0; i < 10; i++) {
+    StatusWithSize copy_result = entry_.Copy(new_address + (i * entry_.size()));
+    ASSERT_EQ(Status::OK, copy_result.status());
+    ASSERT_EQ(kEntry1.size(), copy_result.size());
+  }
+
+  for (int j = 0; j < 10; j++) {
+    Entry entry;
+    FlashPartition::Address read_address = (new_address + (j * entry_.size()));
+    ASSERT_EQ(Status::OK,
+              Entry::Read(partition_, read_address, kFormatsWithSum, &entry));
+
+    EXPECT_EQ(Status::OK, entry.VerifyChecksumInFlash());
+    EXPECT_EQ(kFormatWithSum.magic, entry.magic());
+    EXPECT_EQ(read_address, entry.address());
+    EXPECT_EQ(kTransactionId1 + 6, entry.transaction_id());
+    EXPECT_FALSE(entry.deleted());
+  }
+}
+
+TEST_F(EntryInFlash, DifferentFormat_UpdatedCopy_FailsWithWrongMagic) {
+  ASSERT_EQ(Status::OK, entry_.Update(kFormatWithSum, kTransactionId1 + 6));
+
+  FlashPartition::Address new_address = entry_.size();
+
+  StatusWithSize copy_result = entry_.Copy(new_address);
+  ASSERT_EQ(Status::OK, copy_result.status());
+  ASSERT_EQ(kEntry1.size(), copy_result.size());
+
+  Entry entry;
+  ASSERT_EQ(Status::DATA_LOSS,
+            Entry::Read(partition_, new_address, kFormats, &entry));
+}
+
 TEST_F(EntryInFlash, UpdateAndCopy_WriteError) {
   flash_.InjectWriteError(FlashError::Unconditional(Status::CANCELLED));
 
diff --git a/pw_kvs/key_value_store.cc b/pw_kvs/key_value_store.cc
index 55d030d..d397117 100644
--- a/pw_kvs/key_value_store.cc
+++ b/pw_kvs/key_value_store.cc
@@ -213,6 +213,10 @@
   // initializing last_new_sector_.
   for (EntryMetadata& metadata : entry_cache_) {
     if (metadata.addresses().size() < redundancy()) {
+      DBG("Key 0x%08x missing copies, has %u, needs %u",
+          unsigned(metadata.hash()),
+          unsigned(metadata.addresses().size()),
+          unsigned(redundancy()));
       error_detected_ = true;
     }
     size_t index = 0;
@@ -249,10 +253,11 @@
   sectors_.set_last_new_sector(newest_key);
 
   if (!empty_sector_found) {
+    DBG("No empty sector found");
     error_detected_ = true;
   }
 
-  if (total_corrupt_bytes != 0 || corrupt_entries != 0) {
+  if (error_detected_) {
     WRN("Corruption detected. Found %zu corrupt bytes and %d corrupt entries.",
         total_corrupt_bytes,
         corrupt_entries);
@@ -613,15 +618,9 @@
   // List of addresses for sectors with space for this entry.
   Address* reserved_addresses = entry_cache_.TempReservedAddressesForWrite();
 
-  // Find sectors to write the entry to. This may involve garbage collecting one
-  // or more sectors.
-  for (size_t i = 0; i < redundancy(); i++) {
-    SectorDescriptor* sector;
-    TRY(GetSectorForWrite(&sector, entry_size, span(reserved_addresses, i)));
-
-    DBG("Found space for entry in sector %u", sectors_.Index(sector));
-    reserved_addresses[i] = sectors_.NextWritableAddress(*sector);
-  }
+  // Find addresses to write the entry to. This may involve garbage collecting
+  // one or more sectors.
+  TRY(GetAddressesForWrite(reserved_addresses, entry_size));
 
   // Write the entry at the first address that was found.
   Entry entry = CreateEntry(reserved_addresses[0], key, value, new_state);
@@ -630,7 +629,7 @@
   // After writing the first entry successfully, update the key descriptors.
   // Once a single new the entry is written, the old entries are invalidated.
   EntryMetadata new_metadata =
-      UpdateKeyDescriptor(entry, key, prior_metadata, prior_size);
+      CreateOrUpdateKeyDescriptor(entry, key, prior_metadata, prior_size);
 
   // Write the additional copies of the entry, if redundancy is greater than 1.
   for (size_t i = 1; i < redundancy(); ++i) {
@@ -641,7 +640,7 @@
   return Status::OK;
 }
 
-KeyValueStore::EntryMetadata KeyValueStore::UpdateKeyDescriptor(
+KeyValueStore::EntryMetadata KeyValueStore::CreateOrUpdateKeyDescriptor(
     const Entry& entry,
     string_view key,
     EntryMetadata* prior_metadata,
@@ -651,15 +650,39 @@
     return entry_cache_.AddNew(entry.descriptor(key), entry.address());
   }
 
+  return UpdateKeyDescriptor(
+      entry, entry.address(), prior_metadata, prior_size);
+}
+
+KeyValueStore::EntryMetadata KeyValueStore::UpdateKeyDescriptor(
+    const Entry& entry,
+    Address new_address,
+    EntryMetadata* prior_metadata,
+    size_t prior_size) {
   // Remove valid bytes for the old entry and its copies, which are now stale.
   for (Address address : prior_metadata->addresses()) {
     sectors_.FromAddress(address).RemoveValidBytes(prior_size);
   }
 
-  prior_metadata->Reset(entry.descriptor(key), entry.address());
+  prior_metadata->Reset(entry.descriptor(prior_metadata->hash()), new_address);
   return *prior_metadata;
 }
 
+Status KeyValueStore::GetAddressesForWrite(Address* write_addresses,
+                                           size_t write_size) {
+  for (size_t i = 0; i < redundancy(); i++) {
+    SectorDescriptor* sector;
+    TRY(GetSectorForWrite(&sector, write_size, span(write_addresses, i)));
+    write_addresses[i] = sectors_.NextWritableAddress(*sector);
+
+    DBG("Found space for entry in sector %u at address %u",
+        sectors_.Index(sector),
+        unsigned(write_addresses[i]));
+  }
+
+  return Status::OK;
+}
+
 // Finds a sector to use for writing a new entry to. Does automatic garbage
 // collection if needed and allowed.
 //
@@ -745,15 +768,19 @@
 
 StatusWithSize KeyValueStore::CopyEntryToSector(Entry& entry,
                                                 SectorDescriptor* new_sector,
-                                                Address& new_address) {
-  new_address = sectors_.NextWritableAddress(*new_sector);
+                                                Address new_address) {
   const StatusWithSize result = entry.Copy(new_address);
 
   TRY_WITH_SIZE(MarkSectorCorruptIfNotOk(result.status(), new_sector));
 
   if (options_.verify_on_write) {
-    TRY_WITH_SIZE(
-        MarkSectorCorruptIfNotOk(entry.VerifyChecksumInFlash(), new_sector));
+    Entry new_entry;
+    TRY_WITH_SIZE(MarkSectorCorruptIfNotOk(
+        Entry::Read(partition_, new_address, formats_, &new_entry),
+        new_sector));
+    // TODO: add test that catches doing the verify on the old entry.
+    TRY_WITH_SIZE(MarkSectorCorruptIfNotOk(new_entry.VerifyChecksumInFlash(),
+                                           new_sector));
   }
   // Entry was written successfully; update descriptor's address and the sector
   // descriptors to reflect the new entry.
@@ -780,7 +807,7 @@
   TRY(sectors_.FindSpaceDuringGarbageCollection(
       &new_sector, entry.size(), metadata.addresses(), reserved_addresses));
 
-  Address new_address;
+  Address new_address = sectors_.NextWritableAddress(*new_sector);
   TRY_ASSIGN(const size_t result_size,
              CopyEntryToSector(entry, new_sector, new_address));
   sectors_.FromAddress(address).RemoveValidBytes(result_size);
@@ -801,6 +828,9 @@
     TRY(Repair());
   }
 
+  // Make sure all the entries are on the primary format.
+  UpdateEntriesToPrimaryFormat();
+
   SectorDescriptor* sector = sectors_.last_new();
 
   // TODO: look in to making an iterator method for cycling through sectors
@@ -891,6 +921,52 @@
   return Status::OK;
 }
 
+Status KeyValueStore::UpdateEntriesToPrimaryFormat() {
+  for (EntryMetadata& prior_metadata : entry_cache_) {
+    Entry entry;
+    TRY(ReadEntry(prior_metadata, entry));
+    if (formats_.primary().magic == entry.magic()) {
+      // Ignore entries that are already on the primary format.
+      continue;
+    }
+
+    DBG("Updating entry 0x%08x from old format [0x%08x] to new format "
+        "[0x%08x]",
+        unsigned(prior_metadata.hash()),
+        unsigned(entry.magic()),
+        unsigned(formats_.primary().magic));
+
+    last_transaction_id_ += 1;
+    TRY(entry.Update(formats_.primary(), last_transaction_id_));
+
+    // List of addresses for sectors with space for this entry.
+    Address* reserved_addresses = entry_cache_.TempReservedAddressesForWrite();
+
+    // Find addresses to write the entry to. This may involve garbage collecting
+    // one or more sectors.
+    TRY(GetAddressesForWrite(reserved_addresses, entry.size()));
+
+    TRY(CopyEntryToSector(entry,
+                          &sectors_.FromAddress(reserved_addresses[0]),
+                          reserved_addresses[0]));
+
+    // After writing the first entry successfully, update the key descriptors.
+    // Once a single new the entry is written, the old entries are invalidated.
+    EntryMetadata new_metadata = UpdateKeyDescriptor(
+        entry, reserved_addresses[0], &prior_metadata, entry.size());
+
+    // Write the additional copies of the entry, if redundancy is greater
+    // than 1.
+    for (size_t i = 1; i < redundancy(); ++i) {
+      TRY(CopyEntryToSector(entry,
+                            &sectors_.FromAddress(reserved_addresses[i]),
+                            reserved_addresses[i]));
+      new_metadata.AddNewAddress(reserved_addresses[i]);
+    }
+  }
+  return Status::OK;
+}
+
 // Add any missing redundant entries/copies for a key.
 Status KeyValueStore::AddRedundantEntries(EntryMetadata& metadata) {
   SectorDescriptor* new_sector;
@@ -905,7 +981,7 @@
        i++) {
     TRY(sectors_.FindSpace(&new_sector, entry.size(), metadata.addresses()));
 
-    Address new_address;
+    Address new_address = sectors_.NextWritableAddress(*new_sector);
     TRY(CopyEntryToSector(entry, new_sector, new_address));
 
     metadata.AddNewAddress(new_address);
diff --git a/pw_kvs/key_value_store_binary_format_test.cc b/pw_kvs/key_value_store_binary_format_test.cc
index 43536a4..5ab9c4d 100644
--- a/pw_kvs/key_value_store_binary_format_test.cc
+++ b/pw_kvs/key_value_store_binary_format_test.cc
@@ -581,12 +581,12 @@
 constexpr auto kNoChecksumEntry =
     MakeValidEntry<NoChecksum>(kNoChecksumMagic, 64, "kee", ByteStr("O_o"));
 
-class InitializedMultiMagicKvs : public ::testing::Test {
+class InitializedRedundantMultiMagicKvs : public ::testing::Test {
  protected:
   static constexpr auto kInitialContents =
       AsBytes(kNoChecksumEntry, kEntry1, kAltEntry, kEntry2, kEntry3);
 
-  InitializedMultiMagicKvs()
+  InitializedRedundantMultiMagicKvs()
       : flash_(internal::Entry::kMinAlignmentBytes),
         partition_(&flash_),
         kvs_(&partition_,
@@ -618,7 +618,7 @@
     ASSERT_STREQ(str_value, val);                                      \
   } while (0)
 
-TEST_F(InitializedMultiMagicKvs, AllEntriesArePresent) {
+TEST_F(InitializedRedundantMultiMagicKvs, AllEntriesArePresent) {
   ASSERT_CONTAINS_ENTRY("key1", "value1");
   ASSERT_CONTAINS_ENTRY("k2", "value2");
   ASSERT_CONTAINS_ENTRY("k3y", "value3");
@@ -626,7 +626,7 @@
   ASSERT_CONTAINS_ENTRY("kee", "O_o");
 }
 
-TEST_F(InitializedMultiMagicKvs, RecoversLossOfFirstSector) {
+TEST_F(InitializedRedundantMultiMagicKvs, RecoversLossOfFirstSector) {
   auto stats = kvs_.GetStorageStats();
   EXPECT_EQ(stats.in_use_bytes, (160u * kvs_.redundancy()));
   EXPECT_EQ(stats.reclaimable_bytes, 0u);
@@ -660,7 +660,7 @@
   EXPECT_EQ(stats.missing_redundant_entries_recovered, 10u);
 }
 
-TEST_F(InitializedMultiMagicKvs, RecoversLossOfSecondSector) {
+TEST_F(InitializedRedundantMultiMagicKvs, RecoversLossOfSecondSector) {
   auto stats = kvs_.GetStorageStats();
   EXPECT_EQ(stats.in_use_bytes, (160u * kvs_.redundancy()));
   EXPECT_EQ(stats.reclaimable_bytes, 0u);
@@ -687,7 +687,7 @@
   EXPECT_EQ(stats.missing_redundant_entries_recovered, 10u);
 }
 
-TEST_F(InitializedMultiMagicKvs, SingleReadErrors) {
+TEST_F(InitializedRedundantMultiMagicKvs, SingleReadErrors) {
   // Inject 2 read errors, so the first read attempt fully fails.
   flash_.InjectReadError(FlashError::Unconditional(Status::INTERNAL, 2));
 
@@ -709,7 +709,7 @@
   EXPECT_EQ(stats.missing_redundant_entries_recovered, 5u);
 }
 
-TEST_F(InitializedMultiMagicKvs, SingleWriteError) {
+TEST_F(InitializedRedundantMultiMagicKvs, SingleWriteError) {
   flash_.InjectWriteError(FlashError::Unconditional(Status::INTERNAL, 1, 1));
 
   EXPECT_EQ(Status::INTERNAL, kvs_.Put("new key", ByteStr("abcd?")));
@@ -740,7 +740,7 @@
             kvs_.Get("new key", as_writable_bytes(span(val))).status());
 }
 
-TEST_F(InitializedMultiMagicKvs, DataLossAfterLosingBothCopies) {
+TEST_F(InitializedRedundantMultiMagicKvs, DataLossAfterLosingBothCopies) {
   EXPECT_EQ(Status::OK, partition_.Erase(0, 2));
 
   char val[20] = {};
@@ -765,12 +765,115 @@
   EXPECT_EQ(stats.missing_redundant_entries_recovered, 5u);
 }
 
-class RedundantKvsInitializedSingleCopyData : public ::testing::Test {
+TEST_F(InitializedRedundantMultiMagicKvs, PutNewEntry_UsesFirstFormat) {
+  EXPECT_EQ(Status::OK, kvs_.Put("new key", ByteStr("abcd?")));
+
+  constexpr auto kNewEntry =
+      MakeValidEntry(kMagic, 65, "new key", ByteStr("abcd?"));
+  EXPECT_EQ(0,
+            std::memcmp(kNewEntry.data(),
+                        flash_.buffer().data() + kInitialContents.size(),
+                        kNewEntry.size()));
+  ASSERT_CONTAINS_ENTRY("new key", "abcd?");
+}
+
+TEST_F(InitializedRedundantMultiMagicKvs, PutExistingEntry_UsesFirstFormat) {
+  EXPECT_EQ(Status::OK, kvs_.Put("A Key", ByteStr("New value!")));
+
+  constexpr auto kNewEntry =
+      MakeValidEntry(kMagic, 65, "A Key", ByteStr("New value!"));
+  EXPECT_EQ(0,
+            std::memcmp(kNewEntry.data(),
+                        flash_.buffer().data() + kInitialContents.size(),
+                        kNewEntry.size()));
+  ASSERT_CONTAINS_ENTRY("A Key", "New value!");
+}
+
+#define ASSERT_KVS_CONTAINS_ENTRY(kvs, key, str_value)                \
+  do {                                                                \
+    char val[sizeof(str_value)] = {};                                 \
+    StatusWithSize stat = kvs.Get(key, as_writable_bytes(span(val))); \
+    ASSERT_EQ(Status::OK, stat.status());                             \
+    ASSERT_EQ(sizeof(str_value) - 1, stat.size());                    \
+    ASSERT_STREQ(str_value, val);                                     \
+  } while (0)
+
+TEST_F(InitializedRedundantMultiMagicKvs, UpdateEntryFormat) {
+  ASSERT_EQ(Status::OK, kvs_.FullMaintenance());
+
+  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors, 2, 1> local_kvs(
+      &partition_, {.magic = kMagic, .checksum = &checksum}, kNoGcOptions);
+
+  ASSERT_EQ(Status::OK, local_kvs.Init());
+  EXPECT_EQ(false, local_kvs.error_detected());
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "key1", "value1");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "k2", "value2");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "k3y", "value3");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "A Key", "XD");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "kee", "O_o");
+}
+
+class InitializedMultiMagicKvs : public ::testing::Test {
+ protected:
+  static constexpr auto kInitialContents =
+      AsBytes(kNoChecksumEntry, kEntry1, kAltEntry, kEntry2, kEntry3);
+
+  InitializedMultiMagicKvs()
+      : flash_(internal::Entry::kMinAlignmentBytes),
+        partition_(&flash_),
+        kvs_(&partition_,
+             {{
+                 {.magic = kMagic, .checksum = &checksum},
+                 {.magic = kAltMagic, .checksum = &alt_checksum},
+                 {.magic = kNoChecksumMagic, .checksum = nullptr},
+             }},
+             kRecoveryNoGcOptions) {
+    partition_.Erase();
+    std::memcpy(flash_.buffer().data(),
+                kInitialContents.data(),
+                kInitialContents.size());
+
+    EXPECT_EQ(Status::OK, kvs_.Init());
+  }
+
+  FakeFlashBuffer<512, 4, 3> flash_;
+  FlashPartition partition_;
+  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors, 1, 3> kvs_;
+};
+
+// Similar to test for InitializedRedundantMultiMagicKvs. Doing similar test
+// with different KVS configuration.
+TEST_F(InitializedMultiMagicKvs, AllEntriesArePresent) {
+  ASSERT_CONTAINS_ENTRY("key1", "value1");
+  ASSERT_CONTAINS_ENTRY("k2", "value2");
+  ASSERT_CONTAINS_ENTRY("k3y", "value3");
+  ASSERT_CONTAINS_ENTRY("A Key", "XD");
+  ASSERT_CONTAINS_ENTRY("kee", "O_o");
+}
+
+// Similar to test for InitializedRedundantMultiMagicKvs. Doing similar test
+// with different KVS configuration.
+TEST_F(InitializedMultiMagicKvs, UpdateEntryFormat) {
+  ASSERT_EQ(Status::OK, kvs_.FullMaintenance());
+
+  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors, 1, 1> local_kvs(
+      &partition_, {.magic = kMagic, .checksum = &checksum}, kNoGcOptions);
+
+  ASSERT_EQ(Status::OK, local_kvs.Init());
+  EXPECT_EQ(false, local_kvs.error_detected());
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "key1", "value1");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "k2", "value2");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "k3y", "value3");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "A Key", "XD");
+  ASSERT_KVS_CONTAINS_ENTRY(local_kvs, "kee", "O_o");
+}
+
+class InitializedRedundantLazyRecoveryKvs : public ::testing::Test {
  protected:
   static constexpr auto kInitialContents =
       AsBytes(kEntry1, kEntry2, kEntry3, kEntry4);
 
-  RedundantKvsInitializedSingleCopyData()
+  InitializedRedundantLazyRecoveryKvs()
       : flash_(internal::Entry::kMinAlignmentBytes),
         partition_(&flash_),
         kvs_(&partition_,
@@ -789,7 +892,7 @@
   KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors, 2> kvs_;
 };
 
-TEST_F(RedundantKvsInitializedSingleCopyData, WriteAfterDataLoss) {
+TEST_F(InitializedRedundantLazyRecoveryKvs, WriteAfterDataLoss) {
   EXPECT_EQ(Status::OK, partition_.Erase(0, 4));
 
   char val[20] = {};
@@ -822,8 +925,7 @@
   EXPECT_EQ(stats.missing_redundant_entries_recovered, 4u);
 }
 
-TEST_F(RedundantKvsInitializedSingleCopyData,
-       TwoSectorsCorruptWithGoodEntries) {
+TEST_F(InitializedRedundantLazyRecoveryKvs, TwoSectorsCorruptWithGoodEntries) {
   ASSERT_CONTAINS_ENTRY("key1", "value1");
   ASSERT_CONTAINS_ENTRY("k2", "value2");
   ASSERT_CONTAINS_ENTRY("k3y", "value3");
@@ -858,29 +960,5 @@
   EXPECT_EQ(stats.missing_redundant_entries_recovered, 8u);
 }
 
-TEST_F(InitializedMultiMagicKvs, PutNewEntry_UsesFirstFormat) {
-  EXPECT_EQ(Status::OK, kvs_.Put("new key", ByteStr("abcd?")));
-
-  constexpr auto kNewEntry =
-      MakeValidEntry(kMagic, 65, "new key", ByteStr("abcd?"));
-  EXPECT_EQ(0,
-            std::memcmp(kNewEntry.data(),
-                        flash_.buffer().data() + kInitialContents.size(),
-                        kNewEntry.size()));
-  ASSERT_CONTAINS_ENTRY("new key", "abcd?");
-}
-
-TEST_F(InitializedMultiMagicKvs, PutExistingEntry_UsesFirstFormat) {
-  EXPECT_EQ(Status::OK, kvs_.Put("A Key", ByteStr("New value!")));
-
-  constexpr auto kNewEntry =
-      MakeValidEntry(kMagic, 65, "A Key", ByteStr("New value!"));
-  EXPECT_EQ(0,
-            std::memcmp(kNewEntry.data(),
-                        flash_.buffer().data() + kInitialContents.size(),
-                        kNewEntry.size()));
-  ASSERT_CONTAINS_ENTRY("A Key", "New value!");
-}
-
 }  // namespace
 }  // namespace pw::kvs
diff --git a/pw_kvs/public/pw_kvs/key_value_store.h b/pw_kvs/public/pw_kvs/key_value_store.h
index f5d7d69..08e3761 100644
--- a/pw_kvs/public/pw_kvs/key_value_store.h
+++ b/pw_kvs/public/pw_kvs/key_value_store.h
@@ -403,11 +403,18 @@
                     EntryMetadata* prior_metadata = nullptr,
                     size_t prior_size = 0);
 
-  EntryMetadata UpdateKeyDescriptor(const Entry& new_entry,
-                                    std::string_view key,
+  EntryMetadata CreateOrUpdateKeyDescriptor(const Entry& new_entry,
+                                            std::string_view key,
+                                            EntryMetadata* prior_metadata,
+                                            size_t prior_size);
+
+  EntryMetadata UpdateKeyDescriptor(const Entry& entry,
+                                    Address new_address,
                                     EntryMetadata* prior_metadata,
                                     size_t prior_size);
 
+  Status GetAddressesForWrite(Address* write_addresses, size_t write_size);
+
   Status GetSectorForWrite(SectorDescriptor** sector,
                            size_t entry_size,
                            span<const Address> addresses_to_skip);
@@ -420,7 +427,7 @@
 
   StatusWithSize CopyEntryToSector(Entry& entry,
                                    SectorDescriptor* new_sector,
-                                   Address& new_address);
+                                   Address new_address);
 
   Status RelocateEntry(const EntryMetadata& metadata,
                        KeyValueStore::Address& address,
@@ -437,6 +444,10 @@
   Status GarbageCollectSector(SectorDescriptor& sector_to_gc,
                               span<const Address> addresses_to_skip);
 
+  // Ensure that all entries are on the primary (first) format. Entries that are
+  // not on the primary format are rewritten.
+  Status UpdateEntriesToPrimaryFormat();
+
   Status AddRedundantEntries(EntryMetadata& metadata);
 
   Status RepairCorruptSectors();