Add keys() and remove() methods to SharedData

Signed-off-by: James Mulcahy <[email protected]>

Author: jamesmulcahy

src/shared_data.cc

@@ -56,6 +56,42 @@ WasmResult SharedData::get(std::string_view vm_id, const std::string_view key,
   return WasmResult::NotFound;
 }
 
+WasmResult SharedData::keys(std::string_view vm_id, const std::string_view key_prefix, std::vector<std::string> *result) {
+  std::lock_guard<std::mutex> lock(mutex_);
+  auto map = data_.find(std::string(vm_id));
+  if (map == data_.end()) {
+    return WasmResult::NotFound;
+  }
+
+  for (auto kv : map->second) {
+    if (kv.first.rfind(key_prefix, 0) == 0) {
+      result->push_back(kv.first);
+    }
+  }
+
+  return WasmResult::Ok;
+}
+
+WasmResult SharedData::remove(std::string_view vm_id, std::string_view key, uint32_t cas) {
+  std::lock_guard<std::mutex> lock(mutex_);
+  std::unordered_map<std::string, std::pair<std::string, uint32_t>> *map;
+  auto map_it = data_.find(std::string(vm_id));
+  if (map_it == data_.end()) {
+    return WasmResult::Ok;
+  } else {
+    map = &map_it->second;
+  }
+
+  auto it = map->find(std::string(key));
+  if (it != map->end()) {
+    if (cas && cas != it->second.second) {
+      return WasmResult::CasMismatch;
+    }
+    map->erase(it);
+  }
+  return WasmResult::Ok;
+}
+
 WasmResult SharedData::set(std::string_view vm_id, std::string_view key, std::string_view value,
                            uint32_t cas) {
   std::lock_guard<std::mutex> lock(mutex_);

src/shared_data.h

@@ -25,6 +25,8 @@ class SharedData {
   SharedData(bool register_vm_id_callback = true);
   WasmResult get(std::string_view vm_id, const std::string_view key,
                  std::pair<std::string, uint32_t> *result);
+  WasmResult keys(std::string_view vm_id, const std::string_view key_prefix, std::vector<std::string> *result);
+  WasmResult remove(std::string_view vm_id, const std::string_view key, uint32_t cas);
   WasmResult set(std::string_view vm_id, std::string_view key, std::string_view value,
                  uint32_t cas);
   void deleteByVmId(std::string_view vm_id);

test/shared_data_test.cc

@@ -44,6 +44,25 @@ TEST(SharedData, SingleThread) {
   EXPECT_EQ(WasmResult::Ok, shared_data.get(vm_id, key, &result));
   EXPECT_EQ(value, result.first);
   EXPECT_EQ(result.second, 3);
+
+  std::vector<std::string> keys;
+  EXPECT_EQ(WasmResult::Ok, shared_data.keys(vm_id, "unmatched-prefix", &keys));
+  EXPECT_EQ(0, keys.size());
+
+  EXPECT_EQ(WasmResult::Ok, shared_data.keys(vm_id, "keyyyyy", &keys));
+  EXPECT_EQ(0, keys.size());
+
+  EXPECT_EQ(WasmResult::Ok, shared_data.keys(vm_id, "ke", &keys));
+  EXPECT_EQ(1, keys.size());
+  EXPECT_EQ(key, keys[0]);
+
+  keys.clear();
+  EXPECT_EQ(WasmResult::CasMismatch, shared_data.remove(vm_id, key, 911));
+  EXPECT_EQ(WasmResult::Ok, shared_data.keys(vm_id, "ke", &keys));
+  EXPECT_EQ(1, keys.size());
+
+  EXPECT_EQ(WasmResult::Ok, shared_data.remove(vm_id, key, 0));
+  EXPECT_EQ(WasmResult::NotFound, shared_data.get(vm_id, key, &result));
 }
 
 void incrementData(SharedData *shared_data, std::string_view vm_id, std::string_view key) {