Register the queue cleanup callback to vm_id handle callbacks (proxy-wasm#114)

Signed-off-by: mathetake <[email protected]>

Author: mathetake

src/shared_data.cc

@@ -30,8 +30,10 @@ SharedData &getGlobalSharedData() {
   return *ptr;
 };
 
-SharedData::SharedData() {
-  registerVmIdHandleCallback([this](std::string_view vm_id) { this->deleteByVmId(vm_id); });
+SharedData::SharedData(bool register_vm_id_callback) {
+  if (register_vm_id_callback) {
+    registerVmIdHandleCallback([this](std::string_view vm_id) { this->deleteByVmId(vm_id); });
+  }
 }
 
 void SharedData::deleteByVmId(std::string_view vm_id) {

src/shared_data.h

@@ -22,7 +22,7 @@ namespace proxy_wasm {
 
 class SharedData {
 public:
-  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 set(std::string_view vm_id, std::string_view key, std::string_view value,

src/shared_queue.cc

@@ -21,13 +21,51 @@
 #include <unordered_map>
 #include <unordered_set>
 
+#include "include/proxy-wasm/vm_id_handle.h"
+
 namespace proxy_wasm {
 
 SharedQueue &getGlobalSharedQueue() {
   static auto *ptr = new SharedQueue;
   return *ptr;
 }
 
+SharedQueue::SharedQueue(bool register_vm_id_callback) {
+  if (register_vm_id_callback) {
+    registerVmIdHandleCallback([this](std::string_view vm_id) { this->deleteByVmId(vm_id); });
+  }
+}
+
+void SharedQueue::deleteByVmId(std::string_view vm_id) {
+  std::lock_guard<std::mutex> lock(mutex_);
+  auto queue_keys = vm_queue_keys_.find(std::string(vm_id));
+  if (queue_keys != vm_queue_keys_.end()) {
+    for (auto queue_key : queue_keys->second) {
+      auto token = queue_tokens_.find(queue_key);
+      if (token != queue_tokens_.end()) {
+        queues_.erase(token->second);
+        queue_tokens_.erase(token);
+      }
+    }
+    vm_queue_keys_.erase(queue_keys);
+  }
+}
+
+uint32_t SharedQueue::nextQueueToken() {
+  // TODO(@mathetake): Should we handle the case where the queue overflows, i.e. the number of used
+  // tokens exceeds the max of uint32? If it overflows, the following loop never exits.
+  while (true) {
+    uint32_t token = next_queue_token_++;
+    if (token == 0) {
+      continue; // 0 is an illegal token.
+    }
+
+    if (queues_.find(token) == queues_.end()) {
+      return token;
+    }
+  }
+}
+
 uint32_t SharedQueue::registerQueue(std::string_view vm_id, std::string_view queue_name,
                                     uint32_t context_id, CallOnThreadFunction call_on_thread,
                                     std::string_view vm_key) {
@@ -36,8 +74,18 @@ uint32_t SharedQueue::registerQueue(std::string_view vm_id, std::string_view que
   auto it = queue_tokens_.insert(std::make_pair(key, static_cast<uint32_t>(0)));
   if (it.second) {
     it.first->second = nextQueueToken();
-    queue_token_set_.insert(it.first->second);
+
+    auto vid = std::string(vm_id);
+    QueueKeySet *queue_keys;
+    auto map_it = vm_queue_keys_.find(vid);
+    if (map_it == vm_queue_keys_.end()) {
+      queue_keys = &vm_queue_keys_[vid];
+    } else {
+      queue_keys = &map_it->second;
+    }
+    queue_keys->insert(key);
   }
+
   uint32_t token = it.first->second;
   auto &q = queues_[token];
   q.vm_key = std::string(vm_key);

src/shared_queue.h

@@ -22,25 +22,18 @@ namespace proxy_wasm {
 
 class SharedQueue {
 public:
+  SharedQueue(bool register_vm_id_callback = true);
+
   uint32_t registerQueue(std::string_view vm_id, std::string_view queue_name, uint32_t context_id,
                          CallOnThreadFunction call_on_thread, std::string_view vm_key);
   uint32_t resolveQueue(std::string_view vm_id, std::string_view queue_name);
   WasmResult dequeue(uint32_t token, std::string *data);
   WasmResult enqueue(uint32_t token, std::string_view value);
 
-private:
-  uint32_t nextQueueToken() {
-    while (true) {
-      uint32_t token = next_queue_token_++;
-      if (token == 0) {
-        continue; // 0 is an illegal token.
-      }
-      if (queue_token_set_.find(token) == queue_token_set_.end()) {
-        return token;
-      }
-    }
-  }
+  void deleteByVmId(std::string_view vm_id);
+  uint32_t nextQueueToken();
 
+private:
   struct Queue {
     std::string vm_key;
     uint32_t context_id;
@@ -50,15 +43,22 @@ class SharedQueue {
 
   // TODO: use std::shared_mutex in C++17.
   std::mutex mutex_;
-  std::map<uint32_t, Queue> queues_;
   uint32_t next_queue_token_ = 1;
+
   struct pair_hash {
     template <class T1, class T2> std::size_t operator()(const std::pair<T1, T2> &pair) const {
       return std::hash<T1>()(pair.first) ^ std::hash<T2>()(pair.second);
     }
   };
+
+  using QueueKeySet = std::unordered_set<std::pair<std::string, std::string>, pair_hash>;
+
+  // vm_id -> queue keys
+  std::unordered_map<std::string, QueueKeySet> vm_queue_keys_;
+  // queue key -> token
   std::unordered_map<std::pair<std::string, std::string>, uint32_t, pair_hash> queue_tokens_;
-  std::unordered_set<uint32_t> queue_token_set_;
+  // token -> queue
+  std::unordered_map<uint32_t, Queue> queues_;
 };
 
 SharedQueue &getGlobalSharedQueue();

test/shared_data_test.cc

@@ -23,7 +23,7 @@
 namespace proxy_wasm {
 
 TEST(SharedData, SingleThread) {
-  SharedData shared_data;
+  SharedData shared_data(false);
   std::pair<std::string, uint32_t> result;
   EXPECT_EQ(WasmResult::NotFound, shared_data.get("non-exist", "non-exists", &result));
 
@@ -63,7 +63,7 @@ void incrementData(SharedData *shared_data, std::string_view vm_id, std::string_
 }
 
 TEST(SharedData, Concurrent) {
-  SharedData shared_data;
+  SharedData shared_data(false);
   std::pair<std::string, uint32_t> result;
 
   std::string_view vm_id = "id";
@@ -80,7 +80,7 @@ TEST(SharedData, Concurrent) {
 }
 
 TEST(SharedData, DeleteByVmId) {
-  SharedData shared_data;
+  SharedData shared_data(false);
   std::string_view vm_id = "id";
   std::string_view key = "key";
   std::string_view value;
@@ -91,7 +91,7 @@ TEST(SharedData, DeleteByVmId) {
   EXPECT_EQ(WasmResult::NotFound, shared_data.get(vm_id, key, &result));
 }
 
-TEST(SharedData, VmIdCleanup) {
+TEST(SharedData, VmIdHandleCleanup) {
   SharedData shared_data;
   std::string_view vm_id = "proxy_wasm_shared_data_test";
   auto handle = getVmIdHandle(vm_id);

test/shared_queue_test.cc

@@ -18,16 +18,29 @@
 
 #include "gtest/gtest.h"
 
+#include "include/proxy-wasm/vm_id_handle.h"
+
 namespace proxy_wasm {
 
+TEST(SharedQueue, NextQueueToken) {
+  SharedQueue shared_queue(false);
+  for (auto i = 1; i < 5; i++) {
+    EXPECT_EQ(i, shared_queue.nextQueueToken());
+  }
+  EXPECT_EQ(5, shared_queue.registerQueue("a", "b", 1, nullptr, "c"));
+}
+
 TEST(SharedQueue, SingleThread) {
-  SharedQueue shared_queue;
+  SharedQueue shared_queue(false);
   std::string_view vm_id = "id";
   std::string_view vm_key = "vm_key";
   std::string_view queue_name = "name";
   uint32_t context_id = 1;
 
-  EXPECT_EQ(1, shared_queue.registerQueue(vm_id, queue_name, context_id, nullptr, vm_key));
+  for (auto i = 0; i < 3; i++) {
+    // same token
+    EXPECT_EQ(1, shared_queue.registerQueue(vm_id, queue_name, context_id, nullptr, vm_key));
+  }
   EXPECT_EQ(1, shared_queue.resolveQueue(vm_id, queue_name));
   EXPECT_EQ(0, shared_queue.resolveQueue(vm_id, "non-exist"));
   EXPECT_EQ(0, shared_queue.resolveQueue("non-exist", queue_name));
@@ -69,7 +82,7 @@ void dequeueData(SharedQueue *shared_queue, uint32_t token, size_t *dequeued_cou
 }
 
 TEST(SharedQueue, Concurrent) {
-  SharedQueue shared_queue;
+  SharedQueue shared_queue(false);
   std::string_view vm_id = "id";
   std::string_view vm_key = "vm_key";
   std::string_view queue_name = "name";
@@ -100,4 +113,47 @@ TEST(SharedQueue, Concurrent) {
   EXPECT_EQ(first_cnt + second_cnt, 200);
 }
 
+TEST(SharedQueue, DeleteByVmId) {
+  SharedQueue shared_queue(false);
+  auto vm_id_1 = "id_1";
+  auto vm_id_2 = "id_2";
+  std::string_view vm_key = "vm_key";
+  uint32_t context_id = 1;
+  auto queue_num_per_vm = 3;
+
+  for (auto i = 1; i < queue_num_per_vm; i++) {
+    EXPECT_EQ(i,
+              shared_queue.registerQueue(vm_id_1, std::to_string(i), context_id, nullptr, vm_key));
+    EXPECT_EQ(i, shared_queue.resolveQueue(vm_id_1, std::to_string(i)));
+  }
+
+  for (auto i = queue_num_per_vm; i < 2 * queue_num_per_vm; i++) {
+    EXPECT_EQ(i,
+              shared_queue.registerQueue(vm_id_2, std::to_string(i), context_id, nullptr, vm_key));
+    EXPECT_EQ(i, shared_queue.resolveQueue(vm_id_2, std::to_string(i)));
+  }
+
+  shared_queue.deleteByVmId(vm_id_1);
+  for (auto i = 1; i < queue_num_per_vm; i++) {
+    EXPECT_EQ(0, shared_queue.resolveQueue(vm_id_1, std::to_string(i)));
+  }
+
+  for (auto i = queue_num_per_vm; i < 2 * queue_num_per_vm; i++) {
+    EXPECT_EQ(i, shared_queue.resolveQueue(vm_id_2, std::to_string(i)));
+  }
+}
+
+TEST(SharedQueue, VmIdHandleCleanup) {
+  SharedQueue shared_queue;
+  std::string_view vm_id = "proxy_wasm_shared_queue_test";
+  std::string_view queue_name = "name";
+
+  auto handle = getVmIdHandle(vm_id);
+  EXPECT_EQ(1, shared_queue.registerQueue(vm_id, queue_name, 1, nullptr, "vm_key"));
+  EXPECT_EQ(1, shared_queue.resolveQueue(vm_id, queue_name));
+
+  handle.reset();
+  EXPECT_EQ(0, shared_queue.resolveQueue(vm_id, queue_name));
+}
+
 } // namespace proxy_wasm