From 66f1f5d01f37c159722c1d6def986ac8bf7bfb5a Mon Sep 17 00:00:00 2001
From: kishorenc <kishorenc@gmail.com>
Date: Tue, 14 Jul 2020 20:55:10 +0530
Subject: [PATCH] More efficient store contains.

---
 include/store.h     | 13 +++++++++++++
 test/store_test.cpp | 28 ++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/include/store.h b/include/store.h
index 669bdcb5..ffb96b8a 100644
--- a/include/store.h
+++ b/include/store.h
@@ -109,6 +109,19 @@ public:
 
     bool contains(const std::string& key) const {
         std::string value;
+        bool value_found;
+        bool key_may_exist = db->KeyMayExist(rocksdb::ReadOptions(), key, &value, &value_found);
+
+        // returns false when key definitely does not exist
+        if(!key_may_exist) {
+            return false;
+        }
+
+        if(value_found) {
+            return true;
+        }
+
+        // otherwise, we have try getting the value
         rocksdb::Status status = db->Get(rocksdb::ReadOptions(), key, &value);
         return status.ok() && !status.IsNotFound();
     }
diff --git a/test/store_test.cpp b/test/store_test.cpp
index 7ae388ee..8dadd873 100644
--- a/test/store_test.cpp
+++ b/test/store_test.cpp
@@ -137,4 +137,32 @@ TEST(StoreTest, GetUpdateSinceInvalidIterator) {
     ASSERT_FALSE(updates_op.ok());
     ASSERT_EQ("Invalid iterator. Master's latest sequence number is 4 but updates are requested from sequence number 2. "
                       "The master's WAL entries might have expired (they are kept only for 24 hours).", updates_op.error());
+}
+
+TEST(StoreTest, Contains) {
+    std::string primary_store_path = "/tmp/typesense_test/primary_store_test";
+    LOG(INFO) << "Truncating and creating: " << primary_store_path;
+    system(("rm -rf "+primary_store_path+" && mkdir -p "+primary_store_path).c_str());
+
+    // add some records, flush and try to query
+    Store primary_store(primary_store_path, 0, 0, true);  // disable WAL
+    primary_store.insert("foo1", "bar1");
+    primary_store.insert("foo2", "bar2");
+    primary_store.flush();
+
+    ASSERT_EQ(true, primary_store.contains("foo1"));
+    ASSERT_EQ(true, primary_store.contains("foo2"));
+    ASSERT_EQ(false, primary_store.contains("foo"));
+    ASSERT_EQ(false, primary_store.contains("foo3"));
+
+    // add more records without flushing and query again
+    primary_store.insert("foo3", "bar1");
+    primary_store.insert("foo4", "bar2");
+    primary_store.flush();
+
+    ASSERT_EQ(true, primary_store.contains("foo1"));
+    ASSERT_EQ(true, primary_store.contains("foo3"));
+    ASSERT_EQ(true, primary_store.contains("foo4"));
+    ASSERT_EQ(false, primary_store.contains("foo"));
+    ASSERT_EQ(false, primary_store.contains("foo5"));
 }
\ No newline at end of file