SharedMemoryBlockManager must be thread safe

Author: David Goodwin

src/servers/shared_memory_block_manager.cc

@@ -43,6 +43,8 @@ SharedMemoryBlockManager::CpuCreate(
 {
   *smb = nullptr;
 
+  std::lock_guard<std::mutex> lock(mu_);
+
   if (blocks_.find(name) != blocks_.end()) {
     return TRTSERVER_ErrorNew(
         TRTSERVER_ERROR_ALREADY_EXISTS,
@@ -67,6 +69,8 @@ SharedMemoryBlockManager::GpuCreate(
 {
   *smb = nullptr;
 
+  std::lock_guard<std::mutex> lock(mu_);
+
   if (blocks_.find(name) != blocks_.end()) {
     return TRTSERVER_ErrorNew(
         TRTSERVER_ERROR_ALREADY_EXISTS,
@@ -89,6 +93,8 @@ SharedMemoryBlockManager::Get(
 {
   *smb = nullptr;
 
+  std::lock_guard<std::mutex> lock(mu_);
+
   auto itr = blocks_.find(name);
   if (itr == blocks_.end()) {
     return TRTSERVER_ErrorNew(
@@ -108,6 +114,8 @@ SharedMemoryBlockManager::Find(
 {
   *smb = nullptr;
 
+  std::lock_guard<std::mutex> lock(mu_);
+
   auto itr = blocks_.find(name);
   if (itr != blocks_.end()) {
     *smb = itr->second;
@@ -122,6 +130,8 @@ SharedMemoryBlockManager::Remove(
 {
   *smb = nullptr;
 
+  std::lock_guard<std::mutex> lock(mu_);
+
   auto itr = blocks_.find(name);
   if (itr != blocks_.end()) {
     *smb = itr->second;
@@ -136,6 +146,8 @@ SharedMemoryBlockManager::Clear()
 {
   std::string failed_blocks;
 
+  std::lock_guard<std::mutex> lock(mu_);
+
   auto it = blocks_.begin();
   while (it != blocks_.cend()) {
     TRTSERVER_Error* err = TRTSERVER_SharedMemoryBlockDelete(it->second);

src/servers/shared_memory_block_manager.h

@@ -102,6 +102,7 @@ class SharedMemoryBlockManager {
   TRTSERVER_Error* Clear();
 
  private:
+  std::mutex mu_;
   std::unordered_map<std::string, TRTSERVER_SharedMemoryBlock*> blocks_;
 };