Add support for shm opts serialization

guptask · igchor · commit 14d974e66242 · 2021-11-02T16:12:54.000+01:00
After introducing file segment type, nameToKey_ does not provide
enough information to recover/remove segments on restart.

This commit fixes that by replacing nameToKey_ with nameToOpts_.

Previously, the Key from nameToKey_ map was only used in a single
DCHECK().
diff --git a/cachelib/allocator/CacheAllocator-inl.h b/cachelib/allocator/CacheAllocator-inl.h
@@ -3398,7 +3398,7 @@ bool CacheAllocator<CacheTrait>::stopReaper(std::chrono::seconds timeout) {
 
 template <typename CacheTrait>
 bool CacheAllocator<CacheTrait>::cleanupStrayShmSegments(
-    const std::string& cacheDir, bool posix /*TODO(SHM_FILE): const std::vector<CacheMemoryTierConfig>& config */) {
+  const std::string& cacheDir, bool posix /*TODO(SHM_FILE): const std::vector<CacheMemoryTierConfig>& config */) {
   if (util::getStatIfExists(cacheDir, nullptr) && util::isDir(cacheDir)) {
     try {
       // cache dir exists. clean up only if there are no other processes
diff --git a/cachelib/shm/PosixShmSegment.h b/cachelib/shm/PosixShmSegment.h
@@ -92,13 +92,13 @@ class PosixShmSegment : public ShmBase {
   // @return true if the segment existed. false otherwise
   static bool removeByName(const std::string& name);
 
+  // returns the key type corresponding to the given name.
+  static std::string createKeyForName(const std::string& name) noexcept;
+
  private:
   static int createNewSegment(const std::string& name);
   static int getExisting(const std::string& name, const ShmSegmentOpts& opts);
 
-  // returns the key type corresponding to the given name.
-  static std::string createKeyForName(const std::string& name) noexcept;
-
   // resize the segment
   // @param size  the new size
   // @return none
diff --git a/cachelib/shm/ShmManager.cpp b/cachelib/shm/ShmManager.cpp
@@ -22,6 +22,7 @@
 
 #include <fstream>
 #include <vector>
+#include <string>
 
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wconversion"
@@ -98,7 +99,7 @@ ShmManager::ShmManager(const std::string& dir, bool usePosix)
   // if file exists, init from it if needed.
   const bool reattach = dropSegments ? false : initFromFile();
   if (!reattach) {
-    DCHECK(nameToKey_.empty());
+    DCHECK(nameToOpts_.empty());
   }
   // Lock file for exclusive access
   lockMetadataFile(metaFile);
@@ -109,7 +110,7 @@ ShmManager::ShmManager(const std::string& dir, bool usePosix)
 }
 
 bool ShmManager::initFromFile() {
-  // restore the nameToKey_ map and destroy the contents of the file.
+  // restore the nameToOpts_ map and destroy the contents of the file.
   const std::string fileName = pathName(controlDir_, kMetaDataFile);
   std::ifstream f(fileName);
   SCOPE_EXIT { f.close(); };
@@ -139,9 +140,16 @@ bool ShmManager::initFromFile() {
   }
 
   for (const auto& kv : *object.nameToKeyMap_ref()) {
-    nameToKey_.insert({kv.first, kv.second});
+    if (kv.second.path == "") {
+      PosixSysVSegmentOpts type;
+      type.usePosix = kv.second.usePosix;
+      nameToOpts_.insert({kv.first, type});
+    } else {
+      FileShmSegmentOpts type;
+      type.path = kv.second.path;
+      nameToOpts_.insert({kv.first, type});
+    }
   }
-
   return true;
 }
 
@@ -157,17 +165,28 @@ typename ShmManager::ShutDownRes ShmManager::writeActiveSegmentsToFile() {
     return ShutDownRes::kFileDeleted;
   }
 
-  // write the shmtype, nameToKey_ map to the file.
+  // write the shmtype, nameToOpts_ map to the file.
   DCHECK(metadataStream_);
 
   serialization::ShmManagerObject object;
 
   object.shmVal_ref() = usePosix_ ? static_cast<int>(ShmVal::SHM_POSIX)
                                   : static_cast<int>(ShmVal::SHM_SYS_V);
 
-  for (const auto& kv : nameToKey_) {
+  for (const auto& kv : nameToOpts_) {
     const auto& name = kv.first;
-    const auto& key = kv.second;
+    serialization::ShmTypeObject key;
+    if (const auto* opts = std::get_if<FileShmSegmentOpts>(&kv.second)) {
+      key.path = opts->path;
+    } else {
+      try {
+        const auto& v = std::get<PosixSysVSegmentOpts>(kv.second);
+        key.usePosix = v.usePosix;
+        key.path = "";
+      } catch(std::bad_variant_access&) {
+        throw std::invalid_argument(folly::sformat("Not a valid segment"));
+      }
+    }
     const auto it = segments_.find(name);
     // segment exists and is active.
     if (it != segments_.end() && it->second->isActive()) {
@@ -199,14 +218,14 @@ typename ShmManager::ShutDownRes ShmManager::shutDown() {
 
   // clear our data.
   segments_.clear();
-  nameToKey_.clear();
+  nameToOpts_.clear();
   return ret;
 }
 
 namespace {
 
 bool removeSegByName(ShmTypeOpts typeOpts, const std::string& uniqueName) {
-  if (auto *v = std::get_if<FileShmSegmentOpts>(&typeOpts)) {
+  if (const auto* v = std::get_if<FileShmSegmentOpts>(&typeOpts)) {
     return FileShmSegment::removeByPath(v->path);
   }
 
@@ -258,22 +277,20 @@ void ShmManager::cleanup(const std::string& dir, bool posix) {
 }
 
 void ShmManager::removeAllSegments() {
-  // TODO(SHM_FILE): extend this once we have opts stored in nameToKey_
-  for (const auto& kv : nameToKey_) {
-    removeSegByName(usePosix_, uniqueIdForName(kv.first));
+  for (const auto& kv : nameToOpts_) {
+    removeSegByName(kv.second, uniqueIdForName(kv.first));
   }
-  nameToKey_.clear();
+  nameToOpts_.clear();
 }
 
 void ShmManager::removeUnAttachedSegments() {
-  // TODO(SHM_FILE): extend this once we have opts stored in nameToKey_
-  auto it = nameToKey_.begin();
-  while (it != nameToKey_.end()) {
+  auto it = nameToOpts_.begin();
+  while (it != nameToOpts_.end()) {
     const auto name = it->first;
     // check if the segment is attached.
     if (segments_.find(name) == segments_.end()) { // not attached
-      removeSegByName(usePosix_, uniqueIdForName(name));
-      it = nameToKey_.erase(it);
+      removeSegByName(it->second, uniqueIdForName(name));
+      it = nameToOpts_.erase(it);
     } else {
       ++it;
     }
@@ -292,13 +309,13 @@ ShmAddr ShmManager::createShm(const std::string& shmName,
   removeShm(shmName, opts.typeOpts);
 
   DCHECK(segments_.find(shmName) == segments_.end());
-  DCHECK(nameToKey_.find(shmName) == nameToKey_.end());
+  DCHECK(nameToOpts_.find(shmName) == nameToOpts_.end());
 
-  if (auto *v = std::get_if<PosixSysVSegmentOpts>(&opts.typeOpts)) {
-    if (usePosix_ != v->usePosix)
-      throw std::invalid_argument(
-        folly::sformat("Expected {} but got {} segment",
-        usePosix_ ? "posix" : "SysV", usePosix_ ? "SysV" : "posix"));
+  const auto* v = std::get_if<PosixSysVSegmentOpts>(&opts.typeOpts);
+  if (v && usePosix_ != v->usePosix) {
+    throw std::invalid_argument(
+      folly::sformat("Expected {} but got {} segment",
+      usePosix_ ? "posix" : "SysV", usePosix_ ? "SysV" : "posix"));
   }
 
   std::unique_ptr<ShmSegment> newSeg;
@@ -326,24 +343,32 @@ ShmAddr ShmManager::createShm(const std::string& shmName,
   }
 
   auto ret = newSeg->getCurrentMapping();
-  nameToKey_.emplace(shmName, newSeg->getKeyStr());
+  if (v) {
+    PosixSysVSegmentOpts opts;
+    opts.usePosix = v->usePosix;
+    nameToOpts_.emplace(shmName, opts);
+  } else {
+    FileShmSegmentOpts opts;
+    opts.path = newSeg->getKeyStr();
+    nameToOpts_.emplace(shmName, opts);
+  }
   segments_.emplace(shmName, std::move(newSeg));
   return ret;
 }
 
 void ShmManager::attachNewShm(const std::string& shmName, ShmSegmentOpts opts) {
-  const auto keyIt = nameToKey_.find(shmName);
+  const auto keyIt = nameToOpts_.find(shmName);
   // if key is not known already, there is not much we can do to attach.
-  if (keyIt == nameToKey_.end()) {
+  if (keyIt == nameToOpts_.end()) {
     throw std::invalid_argument(
         folly::sformat("Unable to find any segment with name {}", shmName));
   }
 
-  if (auto *v = std::get_if<PosixSysVSegmentOpts>(&opts.typeOpts)) {
-    if (usePosix_ != v->usePosix)
-      throw std::invalid_argument(
-        folly::sformat("Expected {} but got {} segment",
-        usePosix_ ? "posix" : "SysV", usePosix_ ? "SysV" : "posix"));
+  const auto* v = std::get_if<PosixSysVSegmentOpts>(&opts.typeOpts);
+  if (v && usePosix_ != v->usePosix) {
+    throw std::invalid_argument(
+      folly::sformat("Expected {} but got {} segment",
+      usePosix_ ? "posix" : "SysV", usePosix_ ? "SysV" : "posix"));
   }
 
   // This means the segment exists and we can try to attach it.
@@ -360,7 +385,17 @@ void ShmManager::attachNewShm(const std::string& shmName, ShmSegmentOpts opts) {
         shmName, e.what()));
   }
   DCHECK(segments_.find(shmName) != segments_.end());
-  DCHECK_EQ(segments_[shmName]->getKeyStr(), keyIt->second);
+  if (v) { // If it is a posix shm segment
+    // Comparison unnecessary since getKeyStr() retuns name_from ShmBase
+    // createKeyForShm also returns the same variable.
+  } else { // Else it is a file segment
+    try {
+      auto opts = std::get<FileShmSegmentOpts>(keyIt->second);
+      DCHECK_EQ(segments_[shmName]->getKeyStr(), opts.path);
+    } catch(std::bad_variant_access&) {
+      throw std::invalid_argument(folly::sformat("Not a valid segment"));
+    }
+  }
 }
 
 ShmAddr ShmManager::attachShm(const std::string& shmName,
@@ -403,13 +438,13 @@ bool ShmManager::removeShm(const std::string& shmName, ShmTypeOpts typeOpts) {
         removeSegByName(typeOpts, uniqueIdForName(shmName));
     if (!wasPresent) {
       DCHECK(segments_.end() == segments_.find(shmName));
-      DCHECK(nameToKey_.end() == nameToKey_.find(shmName));
+      DCHECK(nameToOpts_.end() == nameToOpts_.find(shmName));
       return false;
     }
   }
   // not mapped and already removed.
   segments_.erase(shmName);
-  nameToKey_.erase(shmName);
+  nameToOpts_.erase(shmName);
   return true;
 }
 
@@ -424,5 +459,15 @@ ShmSegment& ShmManager::getShmByName(const std::string& shmName) {
   }
 }
 
+ShmTypeOpts& ShmManager::getShmTypeByName(const std::string& shmName) {
+  const auto it = nameToOpts_.find(shmName);
+  if (it != nameToOpts_.end()) {
+    return it->second;
+  } else {
+    throw std::invalid_argument(folly::sformat(
+        "shared memory segment does not exist: name: {}", shmName));
+  }
+}
+
 } // namespace cachelib
 } // namespace facebook
diff --git a/cachelib/shm/ShmManager.h b/cachelib/shm/ShmManager.h
@@ -109,6 +109,14 @@ class ShmManager {
   //         it is returned. Otherwise, it throws std::invalid_argument
   ShmSegment& getShmByName(const std::string& shmName);
 
+  // gets a current segment type by the name that is managed by this
+  // instance. The lifetime of the returned object is same as the
+  // lifetime of this instance.
+  // @param name  Name of the segment
+  // @return If a segment of that name, managed by this instance exists,
+  //         it is returned. Otherwise, it throws std::invalid_argument
+  ShmTypeOpts& getShmTypeByName(const std::string& shmName);
+
   enum class ShutDownRes { kSuccess = 0, kFileDeleted, kFailedWrite };
 
   // persists the metadata information for the current segments managed
@@ -223,8 +231,9 @@ class ShmManager {
   std::unordered_map<std::string, std::unique_ptr<ShmSegment>> segments_{};
 
   // name to key mapping used for reattaching. This is persisted to a
-  // file and used for attaching to the segment.
-  std::unordered_map<std::string, std::string> nameToKey_{};
+  // file using serialization::ShmSegmentVariant and used for attaching
+  // to the segment.
+  std::unordered_map<std::string, ShmTypeOpts> nameToOpts_{};
 
   // file handle for the metadata file. It remains open throughout the lifetime
   // of the object.
diff --git a/cachelib/shm/SysVShmSegment.h b/cachelib/shm/SysVShmSegment.h
@@ -88,10 +88,11 @@ class SysVShmSegment : public ShmBase {
   // @return true if the segment existed. false otherwise
   static bool removeByName(const std::string& name);
 
- private:
   // returns the key identifier for the given name.
   static KeyType createKeyForName(const std::string& name) noexcept;
 
+private:
+
   static int createNewSegment(key_t key,
                               size_t size,
                               const ShmSegmentOpts& opts);
diff --git a/cachelib/shm/shm.thrift b/cachelib/shm/shm.thrift
@@ -16,7 +16,12 @@
 
 namespace cpp2 facebook.cachelib.serialization
 
+struct ShmTypeObject {
+  1: required string path,
+  2: required bool usePosix,
+}
+
 struct ShmManagerObject {
   1: required byte shmVal,
-  3: required map<string, string> nameToKeyMap,
+  3: required map<string, ShmTypeObject> nameToKeyMap,
 }
diff --git a/cachelib/shm/tests/test_shm_manager.cpp b/cachelib/shm/tests/test_shm_manager.cpp
@@ -796,6 +796,9 @@ void ShmManagerTest::testShutDown(bool posix) {
     // destroyed.
     ASSERT_NO_THROW(s.createShm(seg2, seg2Size, nullptr, seg2Opt));
     ASSERT_EQ(s.getShmByName(seg2).getSize(), seg2Size);
+    auto *v = std::get_if<PosixSysVSegmentOpts>(&s.getShmTypeByName(seg2));
+    ASSERT_TRUE(v);
+    ASSERT_EQ(v->usePosix, posix);
 
     ASSERT_TRUE(s.shutDown() == ShutDownRes::kSuccess);
   };
