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

[CCF Archive] Store object type eviction policy submission #3

Closed
kancel wants to merge 382 commits from kancel:ccf-archive-pr2746 into main
pull from: kancel:ccf-archive-pr2746
merge into: mooncake-track:main
Conversation 0 Commits 382 Files Changed 654 +184 -19
4 changed files with 184 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 9153043b86 - Show all commits

2
mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h
View File

@ -104,6 +104,8 @@ class TcpTransport : public Transport {
void startTransfer(Slice *slice);
bool validateAddress(uint64_t addr, uint64_t size) const;
const char *getName() const override { return "tcp"; }
private:

71
mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp
View File

@ -70,12 +70,17 @@ static bool isCudaMemory(void* addr) {
// Forward declaration
class TcpTransport;
using ValidateAddrFn = std::function<bool(uint64_t, uint64_t)>;
// Server-side session: handles one transfer request on a persistent connection
struct ServerSession : public std::enable_shared_from_this<ServerSession> {
explicit ServerSession(std::shared_ptr<tcpsocket> socket)
: socket_(std::move(socket)) {}
explicit ServerSession(std::shared_ptr<tcpsocket> socket,
ValidateAddrFn validate_addr)
: socket_(std::move(socket)),
validate_addr_(std::move(validate_addr)) {}
std::shared_ptr<tcpsocket> socket_;
ValidateAddrFn validate_addr_;
SessionHeader header_;
uint64_t total_transferred_bytes_;
char* local_buffer_;
@ -95,8 +100,6 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
*socket_, asio::buffer(&header_, sizeof(SessionHeader)),
[this, self](const asio::error_code& ec, std::size_t len) {
if (ec || len != sizeof(SessionHeader)) {
// If client closed connection (EOF), this is normal - don't
// log
if (ec.value() != asio::error::eof) {
LOG(WARNING)
<< "ServerSession::readHeader failed. Error: "
@ -104,10 +107,20 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
<< ", bytes read: " << len;
}
session_mutex_.unlock();
return; // Don't continue, socket will be closed
return;
}
local_buffer_ = (char*)(le64toh(header_.addr));
uint64_t size = le64toh(header_.size);
if (validate_addr_ &&
!validate_addr_((uint64_t)local_buffer_, size)) {
LOG(ERROR) << "ServerSession: remote-supplied address 0x"
<< std::hex << (uint64_t)local_buffer_
<< std::dec << " with size " << size
<< " is not within any registered buffer";
session_mutex_.unlock();
return;
}
if (header_.opcode == (uint8_t)TransferRequest::WRITE)
readBody();
else
@ -492,7 +505,8 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
};
struct TcpContext {
TcpContext(short port) : acceptor(io_context) {
TcpContext(short port, ValidateAddrFn validate_addr)
: acceptor(io_context), validate_addr_(std::move(validate_addr)) {
std::error_code ec;
asio::ip::tcp::endpoint endpoint(asio::ip::tcp::v6(), port);
@ -522,11 +536,13 @@ struct TcpContext {
void doAccept() {
acceptor.async_accept([this](asio::error_code ec, tcpsocket socket) {
if (!ec) {
asio::error_code nodelay_ec;
socket.set_option(asio::ip::tcp::no_delay(true), nodelay_ec);
auto socket_ptr =
std::make_shared<tcpsocket>(std::move(socket));
auto session = std::make_shared<ServerSession>(socket_ptr);
session->start(); // Start processing requests on this
// persistent connection
auto session =
std::make_shared<ServerSession>(socket_ptr, validate_addr_);
session->start();
}
doAccept();
});
@ -534,6 +550,7 @@ struct TcpContext {
asio::io_context io_context;
asio::ip::tcp::acceptor acceptor;
ValidateAddrFn validate_addr_;
};
TcpTransport::TcpTransport() : context_(nullptr), running_(false) {
@ -610,7 +627,9 @@ int TcpTransport::install(std::string& local_server_name,
close(sockfd); // the above function has opened a socket
LOG(INFO) << "TcpTransport: listen on port " << tcp_port;
context_ = new TcpContext(tcp_port);
context_ = new TcpContext(tcp_port, [this](uint64_t addr, uint64_t size) {
return validateAddress(addr, size);
});
running_ = true;
thread_ = std::thread(&TcpTransport::worker, this);
return 0;
@ -759,6 +778,7 @@ void TcpTransport::worker() {
LOG(ERROR) << "TcpTransport::worker encountered an exception "
"during doAccept/run: "
<< e.what();
context_->io_context.restart();
}
}
}
@ -774,6 +794,7 @@ std::shared_ptr<asio::ip::tcp::socket> TcpTransport::getConnection(
auto socket_ptr =
std::make_shared<asio::ip::tcp::socket>(context_->io_context);
asio::connect(*socket_ptr, endpoint_iterator);
socket_ptr->set_option(asio::ip::tcp::no_delay(true));
return socket_ptr;
} catch (std::exception& e) {
LOG(ERROR)
@ -826,6 +847,7 @@ std::shared_ptr<asio::ip::tcp::socket> TcpTransport::getConnection(
new_socket =
std::make_shared<asio::ip::tcp::socket>(context_->io_context);
asio::connect(*new_socket, endpoint_iterator);
new_socket->set_option(asio::ip::tcp::no_delay(true));
} catch (std::exception& e) {
LOG(ERROR)
<< "TcpTransport::getConnection failed to create connection to "
@ -900,9 +922,8 @@ void TcpTransport::cleanupIdleConnections() {
for (auto it = connection_pool_.begin(); it != connection_pool_.end();) {
auto& queue = it->second;
// Remove idle connections that exceed timeout
while (!queue.empty()) {
auto& entry = queue.back();
for (auto entry_it = queue.begin(); entry_it != queue.end();) {
auto& entry = *entry_it;
if (!entry->in_use) {
auto idle_duration =
std::chrono::duration_cast<std::chrono::seconds>(
@ -913,16 +934,13 @@ void TcpTransport::cleanupIdleConnections() {
asio::error_code ec;
entry->socket->close(ec);
}
queue.pop_back();
} else {
break;
entry_it = queue.erase(entry_it);
continue;
}
} else {
break;
}
++entry_it;
}
// Remove empty endpoint queues
if (queue.empty()) {
it = connection_pool_.erase(it);
} else {
@ -931,6 +949,21 @@ void TcpTransport::cleanupIdleConnections() {
}
}
bool TcpTransport::validateAddress(uint64_t addr, uint64_t size) const {
if (size == 0) return false;
if (addr + size < addr) return false;
auto desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
if (!desc) return false;
for (const auto& buffer : desc->buffers) {
if (buffer.addr + buffer.length < buffer.addr) continue;
if (buffer.addr <= addr && addr + size <= buffer.addr + buffer.length)
return true;
}
return false;
}
void TcpTransport::startTransfer(Slice* slice) {
auto desc = metadata_->getSegmentDescByID(slice->target_id);
if (!desc) {

4
mooncake-transfer-engine/tests/CMakeLists.txt
View File

@ -65,6 +65,10 @@ target_link_libraries(tcp_transport_test PUBLIC transfer_engine gtest gtest_main
add_test(NAME tcp_transport_test COMMAND tcp_transport_test)
endif()
add_executable(tcp_address_validation_test ${WORKSPACE}/tcp_address_validation_test.cpp)
target_link_libraries(tcp_address_validation_test PUBLIC gtest gtest_main)
add_test(NAME tcp_address_validation_test COMMAND tcp_address_validation_test)
if (USE_MNNVL)
add_executable(nvlink_transport_test ${WORKSPACE}/nvlink_transport_test.cpp)
target_link_libraries(nvlink_transport_test PUBLIC transfer_engine gtest gtest_main )

126
mooncake-transfer-engine/tests/tcp_address_validation_test.cpp Normal file
View File

@ -0,0 +1,126 @@
// Copyright 2024 KVCache.AI
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <gtest/gtest.h>
#include <cstdint>
#include <vector>
namespace {
struct BufferRange {
uint64_t addr;
uint64_t length;
};
// Mirrors TcpTransport::validateAddress logic exactly.
bool validateAddress(uint64_t addr, uint64_t size,
const std::vector<BufferRange>& buffers) {
if (size == 0) return false;
if (addr + size < addr) return false;
for (const auto& buffer : buffers) {
if (buffer.addr + buffer.length < buffer.addr) continue;
if (buffer.addr <= addr && addr + size <= buffer.addr + buffer.length)
return true;
}
return false;
}
class TcpAddressValidationTest : public ::testing::Test {
protected:
std::vector<BufferRange> buffers_;
void SetUp() override {
buffers_ = {
{0x1000, 0x2000}, // [0x1000, 0x3000)
{0x10000, 0x100000}, // [0x10000, 0x110000)
};
}
};
TEST_F(TcpAddressValidationTest, ExactMatch) {
EXPECT_TRUE(validateAddress(0x1000, 0x2000, buffers_));
EXPECT_TRUE(validateAddress(0x10000, 0x100000, buffers_));
}
TEST_F(TcpAddressValidationTest, WithinBounds) {
EXPECT_TRUE(validateAddress(0x1000, 0x100, buffers_));
EXPECT_TRUE(validateAddress(0x1500, 0x500, buffers_));
EXPECT_TRUE(validateAddress(0x2FFF, 1, buffers_));
EXPECT_TRUE(validateAddress(0x50000, 0x1000, buffers_));
}
TEST_F(TcpAddressValidationTest, OutOfBounds) {
EXPECT_FALSE(validateAddress(0x500, 0x100, buffers_));
EXPECT_FALSE(validateAddress(0x3000, 0x100, buffers_));
EXPECT_FALSE(validateAddress(0x5000, 0x1000, buffers_));
EXPECT_FALSE(validateAddress(0x200000, 0x100, buffers_));
}
TEST_F(TcpAddressValidationTest, PartialOverlap) {
EXPECT_FALSE(validateAddress(0x2F00, 0x200, buffers_));
EXPECT_FALSE(validateAddress(0x0F00, 0x200, buffers_));
EXPECT_FALSE(validateAddress(0x10F000, 0x2000, buffers_));
}
TEST_F(TcpAddressValidationTest, ZeroSize) {
EXPECT_FALSE(validateAddress(0x1000, 0, buffers_));
EXPECT_FALSE(validateAddress(0x0, 0, buffers_));
}
TEST_F(TcpAddressValidationTest, IntegerOverflow) {
EXPECT_FALSE(validateAddress(UINT64_MAX, 1, buffers_));
EXPECT_FALSE(validateAddress(UINT64_MAX - 10, 100, buffers_));
EXPECT_FALSE(validateAddress(1, UINT64_MAX, buffers_));
}
TEST_F(TcpAddressValidationTest, EmptyBufferList) {
std::vector<BufferRange> empty;
EXPECT_FALSE(validateAddress(0x1000, 0x100, empty));
}
TEST_F(TcpAddressValidationTest, SingleByteAccess) {
EXPECT_TRUE(validateAddress(0x1000, 1, buffers_));
EXPECT_TRUE(validateAddress(0x2FFF, 1, buffers_));
EXPECT_FALSE(validateAddress(0x3000, 1, buffers_));
EXPECT_FALSE(validateAddress(0x0FFF, 1, buffers_));
}
TEST_F(TcpAddressValidationTest, AdjacentBuffers) {
std::vector<BufferRange> adjacent = {
{0x1000, 0x1000}, // [0x1000, 0x2000)
{0x2000, 0x1000}, // [0x2000, 0x3000)
};
EXPECT_TRUE(validateAddress(0x1000, 0x1000, adjacent));
EXPECT_TRUE(validateAddress(0x2000, 0x1000, adjacent));
// Spanning two buffers should fail
EXPECT_FALSE(validateAddress(0x1800, 0x1000, adjacent));
}
TEST_F(TcpAddressValidationTest, LargeBuffer) {
std::vector<BufferRange> large = {
{0, UINT64_MAX},
};
EXPECT_TRUE(validateAddress(0, 1, large));
EXPECT_TRUE(validateAddress(0, UINT64_MAX, large));
EXPECT_TRUE(validateAddress(UINT64_MAX - 1, 1, large));
}
} // namespace
int main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}