From 5c94db3d130da5c71cbeb79545212106fd8b4aa0 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 20 Aug 2026 12:31:48 +0300 Subject: [PATCH] feat(agent): add PageBroker daemon Signed-off-by: Dan Feigin --- agent/pagebroker/broker.cpp | 285 ++++++++++++++++++ agent/pagebroker/broker.hpp | 42 +++ .../checkpoint_transaction_descriptor.cpp | 33 ++ .../checkpoint_transaction_descriptor.hpp | 23 ++ agent/pagebroker/posix_copy_engine.cpp | 20 ++ agent/pagebroker/posix_copy_engine.hpp | 14 + .../restore_transaction_descriptor.cpp | 19 ++ .../restore_transaction_descriptor.hpp | 18 ++ agent/pagebroker/transfer_engine.cpp | 8 + agent/pagebroker/transfer_engine.hpp | 19 ++ 10 files changed, 481 insertions(+) create mode 100644 agent/pagebroker/broker.cpp create mode 100644 agent/pagebroker/broker.hpp create mode 100644 agent/pagebroker/checkpoint_transaction_descriptor.cpp create mode 100644 agent/pagebroker/checkpoint_transaction_descriptor.hpp create mode 100644 agent/pagebroker/posix_copy_engine.cpp create mode 100644 agent/pagebroker/posix_copy_engine.hpp create mode 100644 agent/pagebroker/restore_transaction_descriptor.cpp create mode 100644 agent/pagebroker/restore_transaction_descriptor.hpp create mode 100644 agent/pagebroker/transfer_engine.cpp create mode 100644 agent/pagebroker/transfer_engine.hpp diff --git a/agent/pagebroker/broker.cpp b/agent/pagebroker/broker.cpp new file mode 100644 index 00000000..6e7c8592 --- /dev/null +++ b/agent/pagebroker/broker.cpp @@ -0,0 +1,285 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "broker.hpp" + +#include + +#include +#include +#include +#include +#include + +#include "posix_copy_engine.hpp" + +namespace snapshot::pagebroker { +namespace fs = std::filesystem; +namespace { + +Response +Reply(const Request& request) +{ + Response response; + response.set_request_id(request.request_id()); + response.set_transaction_id(request.transaction_id()); + return response; +} + +Response +Fail(const Request& request, Failure::Code code, const std::string& message) +{ + auto response = Reply(request); + response.mutable_failure()->set_code(code); + response.mutable_failure()->set_message(message); + return response; +} + +Response +CommitSucceeded(const Request& request) +{ + auto response = Reply(request); + response.mutable_commit_complete(); + return response; +} + +Response +AbortSucceeded(const Request& request) +{ + auto response = Reply(request); + response.mutable_abort_complete(); + return response; +} + +bool +IsSafePathComponent(const std::string& value) +{ + return !value.empty() && value != "." && value != ".." && value.find('/') == std::string::npos && + value.find('\\') == std::string::npos && value.find('\0') == std::string::npos; +} + +bool +IsFilesystemPosix(const StorageBackend& storage, const IOEngine& engine) +{ + return storage.has_filesystem() && !storage.filesystem().directory().empty() && engine.has_posix_copy(); +} + +uintmax_t +TreeSize(const Path& path) +{ + uintmax_t bytes = 0; + for (const auto& entry : fs::recursive_directory_iterator(path)) { + if (entry.is_symlink()) + throw std::runtime_error("checkpoint contains symlink"); + if (entry.is_regular_file()) + bytes += entry.file_size(); + } + return bytes; +} + +bool +HasAvailableSpace(const Path& filesystem, uintmax_t required_bytes) +{ + struct statvfs stat {}; + return statvfs(filesystem.c_str(), &stat) == 0 && uintmax_t(stat.f_bavail) * stat.f_frsize >= required_bytes; +} + +Path +TransactionDirectory(const Path& transaction_root, const std::string& transaction_id) +{ + if (!IsSafePathComponent(transaction_id)) + throw std::runtime_error("invalid transaction path component"); + return transaction_root / transaction_id; +} + +} // namespace + +Broker::Broker(Path staging_root) : staging_root_(fs::weakly_canonical(std::move(staging_root))) +{ + io_engines_.push_back(std::make_unique()); + fs::create_directories(staging_root_ / "restore"); + fs::create_directories(staging_root_ / "checkpoint"); +} + +const TransferEngine& +Broker::Engine(TransferEngineType engine_type) const +{ + for (const auto& candidate : io_engines_) { + if (candidate->type() == engine_type) + return *candidate; + } + throw std::runtime_error("configured I/O engine not found"); +} + +Response +Broker::HandleRequest(const Request& request) +{ + if (!request.has_request_id() || request.request_id().empty() || !request.has_transaction_id() || + !IsSafePathComponent(request.transaction_id())) + return Fail(request, Failure::INVALID_REQUEST, "request and transaction IDs are required"); + + try { + switch (request.command_case()) { + case Request::kStagedRestore: + return Restore(request); + case Request::kPrepareStagedCheckpoint: + return PrepareCheckpoint(request); + case Request::kCommit: + return Commit(request); + case Request::kAbort: + return Abort(request); + default: + return Fail(request, Failure::INVALID_REQUEST, "unsupported operation"); + } + } + catch (const std::exception& error) { + return Fail(request, Failure::STORAGE_ERROR, error.what()); + } +} + +Response +Broker::Restore(const Request& request) +{ + const auto& operation = request.staged_restore(); + if (!IsFilesystemPosix(operation.source(), operation.io_engine())) + return Fail(request, Failure::INVALID_REQUEST, "filesystem storage and POSIX copy are required"); + const auto& engine = Engine(TransferEngineType::POSIX_COPY); + + const Path source(operation.source().filesystem().directory()); + const Path restore_root = staging_root_ / "restore"; + const Path staging_directory = TransactionDirectory(restore_root, request.transaction_id()); + if (!source.is_absolute() || fs::is_symlink(source) || !fs::is_directory(source)) + return Fail(request, Failure::INVALID_REQUEST, "source must be an absolute storage directory"); + if (fs::exists(staging_directory) || transaction_states_.contains(request.transaction_id())) + return Fail(request, Failure::TRANSACTION_CONFLICT, "restore transaction conflicts"); + if (!HasAvailableSpace(staging_root_, TreeSize(source))) + return Fail(request, Failure::INSUFFICIENT_STORAGE, "insufficient tmpfs capacity"); + + try { + transaction_states_.emplace(request.transaction_id(), TransactionState::LIVE); + engine.CopyDirectory(source, staging_directory); + restore_transactions_.emplace(request.transaction_id(), RestoreTransactionDescriptor(staging_directory)); + } + catch (...) { + fs::remove_all(staging_directory); + transaction_states_.erase(request.transaction_id()); + restore_transactions_.erase(request.transaction_id()); + throw; + } + auto response = Reply(request); + response.mutable_staged_restore_directory()->set_image_directory(staging_directory.string()); + return response; +} + +Response +Broker::PrepareCheckpoint(const Request& request) +{ + const auto& operation = request.prepare_staged_checkpoint(); + if (!IsFilesystemPosix(operation.destination(), operation.io_engine())) + return Fail(request, Failure::INVALID_REQUEST, "filesystem storage and POSIX copy are required"); + const auto& engine = Engine(TransferEngineType::POSIX_COPY); + + const Path destination(operation.destination().filesystem().directory()); + const Path checkpoint_root = staging_root_ / "checkpoint"; + const Path staging_directory = TransactionDirectory(checkpoint_root, request.transaction_id()); + if (!destination.is_absolute()) + return Fail(request, Failure::INVALID_REQUEST, "destination must be an absolute storage directory"); + if (fs::exists(staging_directory) || transaction_states_.contains(request.transaction_id())) + return Fail(request, Failure::TRANSACTION_CONFLICT, "checkpoint transaction conflicts"); + try { + transaction_states_.emplace(request.transaction_id(), TransactionState::LIVE); + fs::create_directory(staging_directory); + checkpoint_transactions_.emplace( + request.transaction_id(), + CheckpointTransactionDescriptor(staging_directory, operation.destination(), engine.type())); + } + catch (...) { + fs::remove_all(staging_directory); + transaction_states_.erase(request.transaction_id()); + checkpoint_transactions_.erase(request.transaction_id()); + throw; + } + auto response = Reply(request); + response.mutable_staged_checkpoint_directory()->set_image_directory(staging_directory.string()); + return response; +} + +Response +Broker::Commit(const Request& request) +{ + const auto state = transaction_states_.find(request.transaction_id()); + if (state == transaction_states_.end() || state->second == TransactionState::ABORTED) + return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); + if (state->second == TransactionState::COMMITTED) + return CommitSucceeded(request); + + const auto restore = restore_transactions_.find(request.transaction_id()); + if (restore != restore_transactions_.end()) + return CleanupRestore(request, restore->second); + + const auto checkpoint = checkpoint_transactions_.find(request.transaction_id()); + if (checkpoint == checkpoint_transactions_.end()) + return Fail(request, Failure::INTERNAL_ERROR, "live transaction has no descriptor"); + + return PublishCheckpoint(request, checkpoint->second); +} + +Response +Broker::CleanupRestore(const Request& request, const RestoreTransactionDescriptor& transaction) +{ + fs::remove_all(transaction.staging_directory()); + restore_transactions_.erase(request.transaction_id()); + transaction_states_.at(request.transaction_id()) = TransactionState::COMMITTED; + return CommitSucceeded(request); +} + +Response +Broker::PublishCheckpoint(const Request& request, const CheckpointTransactionDescriptor& transaction) +{ + const Path staging_directory = transaction.staging_directory(); + if (!fs::is_directory(staging_directory)) + return Fail(request, Failure::TRANSACTION_NOT_FOUND, "checkpoint staging directory not found"); + const Path published_directory(transaction.destination_storage().filesystem().directory()); + Path partial = published_directory; + partial += ".pagebroker-partial"; + if (fs::exists(partial)) + return Fail(request, Failure::TRANSACTION_CONFLICT, "checkpoint destination conflicts"); + TreeSize(staging_directory); + + try { + fs::create_directories(published_directory.parent_path()); + Engine(transaction.engine_type()).CopyDirectory(staging_directory, partial); + fs::remove_all(published_directory); + fs::rename(partial, published_directory); + checkpoint_transactions_.erase(request.transaction_id()); + transaction_states_.at(request.transaction_id()) = TransactionState::COMMITTED; + std::error_code cleanup_error; + fs::remove_all(staging_directory, cleanup_error); + } + catch (...) { + fs::remove_all(partial); + throw; + } + return CommitSucceeded(request); +} + +Response +Broker::Abort(const Request& request) +{ + const auto state = transaction_states_.find(request.transaction_id()); + if (state == transaction_states_.end() || state->second == TransactionState::COMMITTED) + return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); + if (state->second == TransactionState::ABORTED) + return AbortSucceeded(request); + + const Path restore_root = staging_root_ / "restore"; + const Path checkpoint_root = staging_root_ / "checkpoint"; + fs::remove_all(TransactionDirectory(restore_root, request.transaction_id())); + fs::remove_all(TransactionDirectory(checkpoint_root, request.transaction_id())); + checkpoint_transactions_.erase(request.transaction_id()); + restore_transactions_.erase(request.transaction_id()); + state->second = TransactionState::ABORTED; + return AbortSucceeded(request); +} + +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/broker.hpp b/agent/pagebroker/broker.hpp new file mode 100644 index 00000000..a7012a8a --- /dev/null +++ b/agent/pagebroker/broker.hpp @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include + +#include "checkpoint_transaction_descriptor.hpp" +#include "pagebroker.pb.h" +#include "restore_transaction_descriptor.hpp" +#include "transfer_engine.hpp" + +namespace snapshot::pagebroker { +class Broker { + public: + explicit Broker(Path staging_root); + Response HandleRequest(const Request& request); + + private: + using Engines = std::vector>; + using CheckpointTransactions = std::unordered_map; + using RestoreTransactions = std::unordered_map; + enum class TransactionState { LIVE, COMMITTED, ABORTED }; + using TransactionStates = std::unordered_map; + + const TransferEngine& Engine(TransferEngineType engine_type) const; + Response Restore(const Request& request); + Response PrepareCheckpoint(const Request& request); + // The Snapshot Agent sends COMMIT after CRIU returns; the provider will send it directly later. + Response Commit(const Request& request); + Response CleanupRestore(const Request& request, const RestoreTransactionDescriptor& transaction); + Response PublishCheckpoint(const Request& request, const CheckpointTransactionDescriptor& transaction); + Response Abort(const Request& request); + Path staging_root_; + Engines io_engines_; + CheckpointTransactions checkpoint_transactions_; + RestoreTransactions restore_transactions_; + TransactionStates transaction_states_; +}; +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/checkpoint_transaction_descriptor.cpp b/agent/pagebroker/checkpoint_transaction_descriptor.cpp new file mode 100644 index 00000000..f457b562 --- /dev/null +++ b/agent/pagebroker/checkpoint_transaction_descriptor.cpp @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "checkpoint_transaction_descriptor.hpp" + +#include + +namespace snapshot::pagebroker { +CheckpointTransactionDescriptor::CheckpointTransactionDescriptor( + Path staging_directory, StorageBackend destination_storage, TransferEngineType engine_type) + : staging_directory_(std::move(staging_directory)), destination_storage_(std::move(destination_storage)), + engine_type_(engine_type) +{ +} + +const Path& +CheckpointTransactionDescriptor::staging_directory() const +{ + return staging_directory_; +} + +const StorageBackend& +CheckpointTransactionDescriptor::destination_storage() const +{ + return destination_storage_; +} + +TransferEngineType +CheckpointTransactionDescriptor::engine_type() const +{ + return engine_type_; +} +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/checkpoint_transaction_descriptor.hpp b/agent/pagebroker/checkpoint_transaction_descriptor.hpp new file mode 100644 index 00000000..a02d2f1d --- /dev/null +++ b/agent/pagebroker/checkpoint_transaction_descriptor.hpp @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "pagebroker.pb.h" +#include "transfer_engine.hpp" + +namespace snapshot::pagebroker { +class CheckpointTransactionDescriptor { + public: + CheckpointTransactionDescriptor(Path staging_directory, StorageBackend destination_storage, TransferEngineType engine_type); + + const Path& staging_directory() const; + const StorageBackend& destination_storage() const; + TransferEngineType engine_type() const; + + private: + Path staging_directory_; + StorageBackend destination_storage_; + TransferEngineType engine_type_; +}; +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/posix_copy_engine.cpp b/agent/pagebroker/posix_copy_engine.cpp new file mode 100644 index 00000000..57b88eac --- /dev/null +++ b/agent/pagebroker/posix_copy_engine.cpp @@ -0,0 +1,20 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "posix_copy_engine.hpp" + +#include + +namespace snapshot::pagebroker { +TransferEngineType +PosixCopyEngine::type() const +{ + return TransferEngineType::POSIX_COPY; +} + +void +PosixCopyEngine::CopyDirectory(const Path& source, const Path& destination) const +{ + std::filesystem::copy(source, destination, std::filesystem::copy_options::recursive); +} +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/posix_copy_engine.hpp b/agent/pagebroker/posix_copy_engine.hpp new file mode 100644 index 00000000..99903493 --- /dev/null +++ b/agent/pagebroker/posix_copy_engine.hpp @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "transfer_engine.hpp" + +namespace snapshot::pagebroker { +class PosixCopyEngine final : public TransferEngine { + public: + TransferEngineType type() const override; + void CopyDirectory(const Path& source, const Path& destination) const override; +}; +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/restore_transaction_descriptor.cpp b/agent/pagebroker/restore_transaction_descriptor.cpp new file mode 100644 index 00000000..629ec87d --- /dev/null +++ b/agent/pagebroker/restore_transaction_descriptor.cpp @@ -0,0 +1,19 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "restore_transaction_descriptor.hpp" + +#include + +namespace snapshot::pagebroker { +RestoreTransactionDescriptor::RestoreTransactionDescriptor(Path staging_directory) + : staging_directory_(std::move(staging_directory)) +{ +} + +const Path& +RestoreTransactionDescriptor::staging_directory() const +{ + return staging_directory_; +} +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/restore_transaction_descriptor.hpp b/agent/pagebroker/restore_transaction_descriptor.hpp new file mode 100644 index 00000000..cfbbb44b --- /dev/null +++ b/agent/pagebroker/restore_transaction_descriptor.hpp @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "transfer_engine.hpp" + +namespace snapshot::pagebroker { +class RestoreTransactionDescriptor { + public: + explicit RestoreTransactionDescriptor(Path staging_directory); + + const Path& staging_directory() const; + + private: + Path staging_directory_; +}; +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/transfer_engine.cpp b/agent/pagebroker/transfer_engine.cpp new file mode 100644 index 00000000..a372a012 --- /dev/null +++ b/agent/pagebroker/transfer_engine.cpp @@ -0,0 +1,8 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "transfer_engine.hpp" + +namespace snapshot::pagebroker { +TransferEngine::~TransferEngine() = default; +} // namespace snapshot::pagebroker diff --git a/agent/pagebroker/transfer_engine.hpp b/agent/pagebroker/transfer_engine.hpp new file mode 100644 index 00000000..711a5a01 --- /dev/null +++ b/agent/pagebroker/transfer_engine.hpp @@ -0,0 +1,19 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +namespace snapshot::pagebroker { +using Path = std::filesystem::path; + +enum class TransferEngineType { POSIX_COPY }; + +class TransferEngine { + public: + virtual ~TransferEngine(); + virtual TransferEngineType type() const = 0; + virtual void CopyDirectory(const Path& source, const Path& destination) const = 0; +}; +} // namespace snapshot::pagebroker