Skip to content

Commit

Permalink
Fix ref_count init bug. Add assert.
Browse files Browse the repository at this point in the history
  • Loading branch information
@small-turtle-1
small-turtle-1 committed Sep 20, 2024
1 parent b8f744c commit e732fdd
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 10 deletions.
39 changes: 33 additions & 6 deletions src/storage/persistence/persistence_manager.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ PersistenceManager::PersistenceManager(const String &workspace, const String &da
current_object_key_ = ObjCreate();
current_object_size_ = 0;
current_object_parts_ = 0;
current_object_ref_count_ = 0;

if (local_data_dir_.empty() || local_data_dir_.back() != '/') {
local_data_dir_ += '/';
Expand Down Expand Up @@ -315,13 +316,32 @@ ObjAddr PersistenceManager::GetObjCache(const String &file_path) {
auto oit = objects_.find(it->second.obj_key_);
if (oit != objects_.end()) {
oit->second.ref_count_++;
LOG_TRACE(fmt::format("GetObjCache object {} ref count {}", it->second.obj_key_, oit->second.ref_count_));
} else {
if (it->second.obj_key_ != current_object_key_) {
String error_message = fmt::format("GetObjCache object {} not found", it->second.obj_key_);
UnrecoverableError(error_message);
return ObjAddr();
}
current_object_ref_count_++;
LOG_TRACE(fmt::format("GetObjCache current object {} ref count {}", it->second.obj_key_, current_object_ref_count_));
}
return it->second;
}

ObjAddr PersistenceManager::GetObjCacheWithoutCnt(const String &local_path) {
String lock_path = RemovePrefix(local_path);
if (lock_path.empty()) {
String error_message = fmt::format("Failed to find local path of {}", local_path);
UnrecoverableError(error_message);
}

std::lock_guard<std::mutex> lock(mtx_);
auto it = local_path_obj_.find(lock_path);
if (it == local_path_obj_.end()) {
String error_message = fmt::format("GetObjCacheWithoutCnt Failed to find object for local path {}", lock_path);
LOG_WARN(error_message);
return ObjAddr();
}
return it->second;
}
Expand Down Expand Up @@ -352,13 +372,15 @@ void PersistenceManager::PutObjCache(const String &file_path) {
UnrecoverableError(fmt::format("PutObjCache object {} ref count is {}", it->second.obj_key_, current_object_ref_count_));
}
current_object_ref_count_--;
LOG_TRACE(fmt::format("PutObjCache current object {} ref count {}", it->second.obj_key_, current_object_ref_count_));
return;
}

if (oit->second.ref_count_ <= 0) {
UnrecoverableError(fmt::format("PutObjCache object {} ref count is {}", it->second.obj_key_, oit->second.ref_count_));
}
oit->second.ref_count_--;
LOG_TRACE(fmt::format("PutObjCache object {} ref count {}", it->second.obj_key_, oit->second.ref_count_));
}

String PersistenceManager::ObjCreate() { return UUID().to_string(); }
Expand Down Expand Up @@ -393,7 +415,7 @@ void PersistenceManager::CurrentObjAppendNoLock(const String &tmp_file_path, Siz
dstFile.close();
}

void PersistenceManager::CleanupNoLock(const ObjAddr &object_addr) {
void PersistenceManager::CleanupNoLock(const ObjAddr &object_addr, bool check_ref_count) {
auto it = objects_.find(object_addr.obj_key_);
if (it == objects_.end()) {
if (object_addr.obj_key_ == current_object_key_) {
Expand All @@ -406,6 +428,13 @@ void PersistenceManager::CleanupNoLock(const ObjAddr &object_addr) {
return;
}
}
if (check_ref_count) {
const ObjStat &stat = it->second;
if (stat.ref_count_ > 0) {
String error_message = fmt::format("CleanupNoLock object {} ref count is {}", object_addr.obj_key_, stat.ref_count_);
UnrecoverableError(error_message);
}
}
Range orig_range(object_addr.part_offset_, object_addr.part_offset_ + object_addr.part_size_);
Range range(orig_range);
auto inst_it = it->second.deleted_ranges_.lower_bound(range);
Expand Down Expand Up @@ -544,7 +573,7 @@ void PersistenceManager::Cleanup(const String &file_path) {
LOG_WARN(error_message);
return;
}
CleanupNoLock(it->second);
CleanupNoLock(it->second, true);
LOG_TRACE(fmt::format("Deleted mapping from local path {} to ObjAddr({}, {}, {})",
local_path,
it->second.obj_key_,
Expand Down Expand Up @@ -622,7 +651,7 @@ void AddrSerializer::Initialize(PersistenceManager *persistence_manager, const V
}
for (const String &path : path) {
paths_.push_back(path);
ObjAddr obj_addr = persistence_manager->GetObjCache(path);
ObjAddr obj_addr = persistence_manager->GetObjCacheWithoutCnt(path);
obj_addrs_.push_back(obj_addr);
if (!obj_addr.Valid()) {
// In ImportWal, version file is not flushed here, set before write wal
Expand All @@ -631,7 +660,6 @@ void AddrSerializer::Initialize(PersistenceManager *persistence_manager, const V
} else {
ObjStat obj_stat = persistence_manager->GetObjStatByObjAddr(obj_addr);
obj_stats_.push_back(obj_stat);
persistence_manager->PutObjCache(path);
}
}
}
Expand All @@ -645,15 +673,14 @@ void AddrSerializer::InitializeValid(PersistenceManager *persistence_manager) {
continue;
}

ObjAddr obj_addr = persistence_manager->GetObjCache(paths_[i]);
ObjAddr obj_addr = persistence_manager->GetObjCacheWithoutCnt(paths_[i]);

obj_addrs_[i] = obj_addr;
if (!obj_addr.Valid()) {
UnrecoverableError(fmt::format("Invalid object address for path {}", paths_[i]));
} else {
ObjStat obj_stat = persistence_manager->GetObjStatByObjAddr(obj_addr);
obj_stats_[i] = obj_stat;
persistence_manager->PutObjCache(paths_[i]);
}
}
}
Expand Down
10 changes: 6 additions & 4 deletions src/storage/persistence/persistence_manager.cppm
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,8 @@ public:
// Download the whole object from object store if it's not in cache. Increase refcount and return the cached object file path.
ObjAddr GetObjCache(const String &local_path);

ObjAddr GetObjCacheWithoutCnt(const String &local_path);

void PutObjCache(const String &file_path);

void Cleanup(const String &file_path);
Expand Down Expand Up @@ -118,7 +120,7 @@ private:
void CurrentObjFinalizeNoLock();

// Cleanup
void CleanupNoLock(const ObjAddr &object_addr);
void CleanupNoLock(const ObjAddr &object_addr, bool check_ref_count = false);

String RemovePrefix(const String &path);

Expand All @@ -139,9 +141,9 @@ private:
HashMap<String, ObjAddr> local_path_obj_; // local_file_path -> ObjAddr
// Current unsealed object key
String current_object_key_;
SizeT current_object_size_;
SizeT current_object_parts_;
SizeT current_object_ref_count_;
SizeT current_object_size_ = 0;
SizeT current_object_parts_ = 0;
SizeT current_object_ref_count_ = 0;

friend struct AddrSerializer;
};
Expand Down

0 comments on commit e732fdd

Please sign in to comment.