mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[TransferEngine] implement closeLocalSegment by segment name (#286)

This commit is contained in:
Teng Ma 2025-04-25 17:35:16 +08:00 committed by GitHub
parent e6d7c123c5
commit 1519ce60ad
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 34 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mooncake-transfer-engine/include/transfer_engine.h
View File

@ -89,6 +89,8 @@ class TransferEngine {
int closeSegment(SegmentHandle handle);
int removeLocalSegment(const std::string &segment_name);
int registerLocalMemory(void *addr, size_t length,
const std::string &location = kWildcardLocation,
bool remote_accessible = true,

2
mooncake-transfer-engine/include/transfer_engine_c.h
View File

@ -112,6 +112,8 @@ segment_id_t openSegmentNoCache(transfer_engine_t engine, const char *segment_na
int closeSegment(transfer_engine_t engine, segment_id_t segment_id);
int removeLocalSegment(transfer_engine_t engine, const char *segment_name);
void destroyTransferEngine(transfer_engine_t engine);
int registerLocalMemory(transfer_engine_t engine, void *addr, size_t length,

2
mooncake-transfer-engine/include/transfer_metadata.h
View File

@ -120,6 +120,8 @@ class TransferMetadata {
int addLocalSegment(SegmentID segment_id, const std::string &segment_name,
std::shared_ptr<SegmentDesc> &&desc);
int removeLocalSegment(const std::string &segment_name);
int addRpcMetaEntry(const std::string &server_name, RpcMetaDesc &desc);

9
mooncake-transfer-engine/src/transfer_engine.cpp
View File

@ -199,6 +199,15 @@ Transport::SegmentHandle TransferEngine::openSegment(
int TransferEngine::closeSegment(Transport::SegmentHandle handle) { return 0; }
int TransferEngine::removeLocalSegment(const std::string &segment_name) {
if (segment_name.empty()) return ERR_INVALID_ARGUMENT;
std::string trimmed_segment_name = segment_name;
while (!trimmed_segment_name.empty() && trimmed_segment_name[0] == '/')
trimmed_segment_name.erase(0, 1);
if (trimmed_segment_name.empty()) return ERR_INVALID_ARGUMENT;
return metadata_->removeLocalSegment(trimmed_segment_name);
}
bool TransferEngine::checkOverlap(void *addr, uint64_t length) {
std::shared_lock<std::shared_mutex> lock(mutex_);
for (auto &local_memory_region : local_memory_regions_) {

5
mooncake-transfer-engine/src/transfer_engine_c.cpp
View File

@ -77,6 +77,11 @@ int closeSegment(transfer_engine_t engine, segment_id_t segment_id) {
return native->closeSegment(segment_id);
}
int removeLocalSegment(transfer_engine_t engine, const char *segment_name) {
TransferEngine *native = (TransferEngine *)engine;
return native->removeLocalSegment(segment_name);
}
int registerLocalMemory(transfer_engine_t engine, void *addr, size_t length,
const char *location, int remote_accessible) {
TransferEngine *native = (TransferEngine *)engine;

10
mooncake-transfer-engine/src/transfer_metadata.cpp
View File

@ -380,6 +380,16 @@ int TransferMetadata::addLocalSegment(SegmentID segment_id,
return 0;
}
int TransferMetadata::removeLocalSegment(const std::string &segment_name) {
RWSpinlock::WriteGuard guard(segment_lock_);
if (segment_name_to_id_map_.count(segment_name)) {
int segment_id = segment_name_to_id_map_[segment_name];
segment_name_to_id_map_.erase(segment_name);
segment_id_to_desc_map_.erase(segment_id);
}
return 0;
}
int TransferMetadata::addLocalMemoryBuffer(const BufferDesc &buffer_desc,
bool update_metadata) {
{

4
mooncake-transfer-engine/tests/transfer_metadata_test.cpp
View File

@ -75,6 +75,8 @@ TEST_F(TransferMetadataTest, LocalSegmentTest) {
ASSERT_EQ(des, segment_des);
auto id = metadata_client->getSegmentID(segment_name);
ASSERT_EQ(id, segment_id);
re = metadata_client->removeLocalSegment(segment_name);
ASSERT_EQ(re, 0);
}
// add and remove LocalMemoryBufferMeta
@ -101,6 +103,8 @@ TEST_F(TransferMetadataTest, LocalMemoryBufferTest) {
re = metadata_client->removeLocalMemoryBuffer((void*)addr, false);
ASSERT_EQ(re, 0);
}
re = metadata_client->removeLocalSegment("test_local_segment");
ASSERT_EQ(re, 0);
}
// add, get and remove RPCMetaEntryMeta