Properly handle case where we delete a key that doesn't exist while a write batch is active

src/storage/rocksdb.cpp

@@ -100,14 +100,14 @@ bool RocksDBStorageProvider::erase(const char *key, size_t cchKey)
 {
     rocksdb::Status status;
     std::unique_lock<fastlock> l(m_lock);
+    if (!FKeyExists(key, cchKey))
+        return false;
     if (m_spbatch != nullptr)
     {
         status = m_spbatch->Delete(m_spcolfamily.get(), rocksdb::Slice(key, cchKey));
     }
     else
     {
-        if (!FKeyExists(key, cchKey))
-            return false;
         status = m_spdb->Delete(WriteOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cchKey));
     }
     if (status.ok())
@@ -202,12 +202,12 @@ rocksdb::WriteOptions RocksDBStorageProvider::WriteOptions() const
 void RocksDBStorageProvider::beginWriteBatch()
 {
     m_lock.lock();
-    m_spbatch = std::make_unique<rocksdb::WriteBatch>();
+    m_spbatch = std::make_unique<rocksdb::WriteBatchWithIndex>();
 }
 
 void RocksDBStorageProvider::endWriteBatch()
 {
-    m_spdb->Write(WriteOptions(), m_spbatch.get());
+    m_spdb->Write(WriteOptions(), m_spbatch.get()->GetWriteBatch());
     m_spbatch = nullptr;
     m_lock.unlock();
 }
@@ -230,5 +230,7 @@ void RocksDBStorageProvider::flush()
 bool RocksDBStorageProvider::FKeyExists(const char *key, size_t cch) const
 {
     rocksdb::PinnableSlice slice;
+    if (m_spbatch)
+        return m_spbatch->GetFromBatchAndDB(m_spdb.get(), ReadOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cch), &slice).ok();
     return m_spdb->Get(ReadOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cch), &slice).ok();
 }

src/storage/rocksdb.h

@@ -3,6 +3,7 @@
 #include <memory>
 #include "../IStorage.h"
 #include <rocksdb/db.h>
+#include <rocksdb/utilities/write_batch_with_index.h>
 #include "../fastlock.h"
 
 #define INTERNAL_KEY_PREFIX "\x00\x04\x03\x00\x05\x02\x04"
@@ -15,7 +16,7 @@ class RocksDBStorageProvider : public IStorage
 {
     RocksDBStorageFactory *m_pfactory;
     std::shared_ptr<rocksdb::DB> m_spdb;    // Note: This must be first so it is deleted last
-    std::unique_ptr<rocksdb::WriteBatch> m_spbatch;
+    std::unique_ptr<rocksdb::WriteBatchWithIndex> m_spbatch;
     const rocksdb::Snapshot *m_psnapshot = nullptr;
     std::shared_ptr<rocksdb::ColumnFamilyHandle> m_spcolfamily;
     rocksdb::ReadOptions m_readOptionsTemplate;