diff --git a/src/storage/rocksdb.cpp b/src/storage/rocksdb.cpp index d8acf9f51..e592c74c5 100644 --- a/src/storage/rocksdb.cpp +++ b/src/storage/rocksdb.cpp @@ -100,14 +100,14 @@ bool RocksDBStorageProvider::erase(const char *key, size_t cchKey) { rocksdb::Status status; std::unique_lock 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(); + m_spbatch = std::make_unique(); } 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(); } \ No newline at end of file diff --git a/src/storage/rocksdb.h b/src/storage/rocksdb.h index a35fea056..10c54606c 100644 --- a/src/storage/rocksdb.h +++ b/src/storage/rocksdb.h @@ -3,6 +3,7 @@ #include #include "../IStorage.h" #include +#include #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 m_spdb; // Note: This must be first so it is deleted last - std::unique_ptr m_spbatch; + std::unique_ptr m_spbatch; const rocksdb::Snapshot *m_psnapshot = nullptr; std::shared_ptr m_spcolfamily; rocksdb::ReadOptions m_readOptionsTemplate;