From 232cc7c56674742c1ea2d4b966f156a1433dceba Mon Sep 17 00:00:00 2001
From: John Sully <john@csquare.ca>
Date: Sun, 22 Dec 2019 17:30:15 -0500
Subject: [PATCH] Fix issue where DELETE always returns 1 when using FLASH

Former-commit-id: 0574db59fdc740255388ae77d2ece8622898bb49
---
 src/IStorage.h          |  2 +-
 src/db.cpp              |  2 +-
 src/storage/rocksdb.cpp | 35 +++++++++++++++++++++++++++--------
 src/storage/rocksdb.h   |  2 +-
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/IStorage.h b/src/IStorage.h
index cdb8336f4..5bf583804 100644
--- a/src/IStorage.h
+++ b/src/IStorage.h
@@ -22,7 +22,7 @@ public:
     virtual void retrieve(const char *key, size_t cchKey, callbackSingle fn) const = 0;
     virtual size_t clear() = 0;
     virtual bool enumerate(callback fn) const = 0;
-    virtual size_t count() const = 0;
+    virtual size_t count(bool fStrict) const = 0;
 
     virtual void beginWriteBatch() {} // NOP
     virtual void endWriteBatch() {} // NOP
diff --git a/src/db.cpp b/src/db.cpp
index dc83cff91..f334929dd 100644
--- a/src/db.cpp
+++ b/src/db.cpp
@@ -2195,7 +2195,7 @@ dict_iter redisDbPersistentData::random()
 size_t redisDbPersistentData::size() const 
 { 
     if (m_spstorage != nullptr)
-        return m_spstorage->count();
+        return m_spstorage->count(false);
     
     return dictSize(m_pdict) 
         + (m_pdbSnapshot ? (m_pdbSnapshot->size() - dictSize(m_pdictTombstone)) : 0); 
diff --git a/src/storage/rocksdb.cpp b/src/storage/rocksdb.cpp
index 4addc26b1..cb70d6c83 100644
--- a/src/storage/rocksdb.cpp
+++ b/src/storage/rocksdb.cpp
@@ -24,9 +24,16 @@ bool RocksDBStorageProvider::erase(const char *key, size_t cchKey)
 {
     rocksdb::Status status;
     if (m_spbatch != nullptr)
+    {
         status = m_spbatch->Delete(m_spcolfamily.get(), rocksdb::Slice(key, cchKey));
+    }
     else
+    {
+        std::string strT;
+        if (!m_spdb->KeyMayExist(ReadOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cchKey), &strT))
+            return false;
         status = m_spdb->Delete(WriteOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cchKey));
+    }
     return status.ok();
 }
 
@@ -40,7 +47,7 @@ void RocksDBStorageProvider::retrieve(const char *key, size_t cchKey, callbackSi
 
 size_t RocksDBStorageProvider::clear()
 {
-    size_t celem = count();
+    size_t celem = count(false);
     auto status = m_spdb->DropColumnFamily(m_spcolfamily.get());
     auto strName = m_spcolfamily->GetName();
 
@@ -53,14 +60,26 @@ size_t RocksDBStorageProvider::clear()
     return celem;
 }
 
-size_t RocksDBStorageProvider::count() const
+size_t RocksDBStorageProvider::count(bool fStrict) const
 {
-    std::string strelem;
-    if (!m_spdb->GetProperty(m_spcolfamily.get(), rocksdb::DB::Properties::kEstimateNumKeys, &strelem))
-        throw "Failed to get database size";
-    std::stringstream sstream(strelem);
-    size_t count;
-    sstream >> count;
+    size_t count = 0;
+
+    if (fStrict)
+    {
+        std::unique_ptr<rocksdb::Iterator> it = std::unique_ptr<rocksdb::Iterator>(m_spdb->NewIterator(ReadOptions(), m_spcolfamily.get()));
+        for (it->SeekToFirst(); it->Valid(); it->Next()) {
+            ++count;
+        }
+    }
+    else
+    {
+        std::string strelem;
+        if (!m_spdb->GetProperty(m_spcolfamily.get(), rocksdb::DB::Properties::kEstimateNumKeys, &strelem))
+            throw "Failed to get database size";
+        std::stringstream sstream(strelem);
+        size_t count;
+        sstream >> count;
+    }
     return count;
 }
 
diff --git a/src/storage/rocksdb.h b/src/storage/rocksdb.h
index 84b0bbf31..5ab3a468f 100644
--- a/src/storage/rocksdb.h
+++ b/src/storage/rocksdb.h
@@ -29,7 +29,7 @@ public:
 
     virtual void flush() override;
 
-    size_t count() const;
+    size_t count(bool fStrict) const;
 
 protected: