pw_kvs: Fix bug related to relocation

Fix bug when writing an existing key that has the old entry relocated as
part of a garbage collection that is triggered by the write of the new
key.

Change-Id: I7f0f1e4cb782c629b102dc3b31e022592c2f13e5
diff --git a/pw_kvs/key_value_store.cc b/pw_kvs/key_value_store.cc
index 4a686d8..0c8621e 100644
--- a/pw_kvs/key_value_store.cc
+++ b/pw_kvs/key_value_store.cc
@@ -285,8 +285,9 @@
 
   KeyDescriptor* key_descriptor;
   if (FindKeyDescriptor(key, &key_descriptor).ok()) {
-    DBG("Writing over existing entry for key 0x%08" PRIx32,
-        key_descriptor->key_hash);
+    DBG("Writing over existing entry for key 0x%08" PRIx32 " in sector %zu",
+        key_descriptor->key_hash,
+        SectorIndex(SectorFromAddress(key_descriptor->address)));
     return WriteEntryForExistingKey(
         key_descriptor, KeyDescriptor::kValid, key, value);
   }
@@ -304,7 +305,9 @@
     return Status::NOT_FOUND;
   }
 
-  DBG("Writing tombstone for key 0x%08" PRIx32, key_descriptor->key_hash);
+  DBG("Writing tombstone for existing key 0x%08" PRIx32 " in sector %zu",
+      key_descriptor->key_hash,
+      SectorIndex(SectorFromAddress(key_descriptor->address)));
   return WriteEntryForExistingKey(
       key_descriptor, KeyDescriptor::kDeleted, key, {});
 }
@@ -469,16 +472,25 @@
   // Find the original entry and sector to update the sector's valid_bytes.
   EntryHeader original_entry;
   TRY(ReadEntryHeader(key_descriptor->address, &original_entry));
-  SectorDescriptor& old_sector = SectorFromAddress(key_descriptor->address);
+  SectorDescriptor* old_sector = SectorFromAddress(key_descriptor->address);
 
   SectorDescriptor* sector;
   TRY(FindOrRecoverSectorWithSpace(
       &sector, EntryHeader::size(partition_.alignment_bytes(), key, value)));
-
   DBG("Writing existing entry; found sector: %zu", SectorIndex(sector));
+
+  if (old_sector != SectorFromAddress(key_descriptor->address)) {
+    DBG("Sector for old entry (size %zu) was garbage collected. Old entry "
+        "relocated to sector %zu",
+        original_entry.size(),
+        SectorIndex(SectorFromAddress(key_descriptor->address)));
+
+    old_sector = SectorFromAddress(key_descriptor->address);
+  }
+
   TRY(AppendEntry(sector, key_descriptor, key, value, new_state));
 
-  old_sector.RemoveValidBytes(original_entry.size());
+  old_sector->RemoveValidBytes(original_entry.size());
   return Status::OK;
 }
 
@@ -538,16 +550,16 @@
   TRY(header.VerifyChecksum(
       entry_header_format_.checksum, key, as_bytes(value)));
 
-  SectorDescriptor& old_sector = SectorFromAddress(key_descriptor.address);
+  SectorDescriptor* old_sector = SectorFromAddress(key_descriptor.address);
 
   // Find a new sector for the entry and write it to the new location.
   SectorDescriptor* new_sector;
-  TRY(FindSectorWithSpace(&new_sector, header.size(), &old_sector, true));
+  TRY(FindSectorWithSpace(&new_sector, header.size(), old_sector, true));
   TRY(AppendEntry(new_sector, &key_descriptor, key, as_bytes(value)));
 
   // Do the valid bytes accounting for the sector the entry was relocated out
   // of.
-  old_sector.RemoveValidBytes(header.size());
+  old_sector->RemoveValidBytes(header.size());
 
   return Status::OK;
 }
@@ -683,6 +695,7 @@
 
   // Step 1: Find the sector to garbage collect
   SectorDescriptor* sector_to_gc = FindSectorToGarbageCollect();
+  LogSectors();
 
   if (sector_to_gc == nullptr) {
     return Status::RESOURCE_EXHAUSTED;
@@ -743,7 +756,7 @@
       unsigned(header.key_version()));
 
   Address address = NextWritableAddress(sector);
-  DBG("Appending to address: %zx", size_t(address));
+  DBG("Appending to address: %#zx", size_t(address));
 
   // Handles writing multiple concatenated buffers, while breaking up the writes
   // into alignment-sized blocks.
@@ -852,4 +865,16 @@
   }
 }
 
+void KeyValueStore::SectorDescriptor::RemoveValidBytes(size_t size) {
+  // TODO: add safety check for valid_bytes > size.
+  if (size > valid_bytes) {
+    CRT("!!!!!!!!!!!!!!!");
+    CRT("Remove too many valid bytes!!! remove %zu, only have %hu",
+        size,
+        valid_bytes);
+    valid_bytes = size;
+  }
+  valid_bytes -= size;
+}
+
 }  // 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 3b07b72c8..d8798fc 100644
--- a/pw_kvs/public/pw_kvs/key_value_store.h
+++ b/pw_kvs/public/pw_kvs/key_value_store.h
@@ -252,10 +252,7 @@
       tail_free_bytes -= size;
     }
 
-    void RemoveValidBytes(size_t size) {
-      // TODO: add safety check for valid_bytes > size.
-      valid_bytes -= size;
-    }
+    void RemoveValidBytes(size_t size);
   };
 
   static uint32_t HashKey(std::string_view string);
@@ -350,11 +347,11 @@
     return SectorIndex(sector) * partition_.sector_size_bytes();
   }
 
-  SectorDescriptor& SectorFromAddress(Address address) {
+  SectorDescriptor* SectorFromAddress(Address address) {
     const size_t index = address / partition_.sector_size_bytes();
     // TODO: Add boundary checking once asserts are supported.
     // DCHECK_LT(index, sector_map_size_);
-    return sector_map_[index];
+    return &sector_map_[index];
   }
 
   void LogSectors(void);