From 834095424a709b648404c898ee2097c2c1ec565d Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 11:42:20 +0300 Subject: [PATCH 01/10] feat(snapshot): allow concurrent PageBroker requests --- deploy/snapshot/internal/types/config.go | 3 + deploy/snapshot/internal/types/config_test.go | 9 ++ deploy/snapshot/pagebroker/Makefile | 6 +- deploy/snapshot/pagebroker/broker.cpp | 103 ++++++++++++------ deploy/snapshot/pagebroker/broker.hpp | 16 +-- .../pagebroker/connection_tracker.cpp | 23 ++++ .../pagebroker/connection_tracker.hpp | 16 +++ deploy/snapshot/pagebroker/daemon.cpp | 47 +++++++- deploy/snapshot/pagebroker/daemon_test.cpp | 20 ++++ deploy/snapshot/pagebroker/transaction.cpp | 53 +++++++++ deploy/snapshot/pagebroker/transaction.hpp | 31 ++++++ 11 files changed, 281 insertions(+), 46 deletions(-) create mode 100644 deploy/snapshot/pagebroker/connection_tracker.cpp create mode 100644 deploy/snapshot/pagebroker/connection_tracker.hpp create mode 100644 deploy/snapshot/pagebroker/transaction.cpp create mode 100644 deploy/snapshot/pagebroker/transaction.hpp diff --git a/deploy/snapshot/internal/types/config.go b/deploy/snapshot/internal/types/config.go index f76dccbf80c6..d517ca412e38 100644 --- a/deploy/snapshot/internal/types/config.go +++ b/deploy/snapshot/internal/types/config.go @@ -67,6 +67,9 @@ func (c *AgentConfig) Validate() error { } } c.Storage.AccessMode = accessMode + if c.PageBroker.Enabled && strings.TrimSpace(c.PageBroker.ControlSocketPath) == "" { + return &ConfigError{Field: "pageBroker.controlSocketPath", Message: "pageBroker.controlSocketPath is required when PageBroker is enabled"} + } if c.CRIU.TcpClose && c.CRIU.TcpEstablished { return &ConfigError{ Field: "criu", diff --git a/deploy/snapshot/internal/types/config_test.go b/deploy/snapshot/internal/types/config_test.go index 962207da0535..f247acaee9b4 100644 --- a/deploy/snapshot/internal/types/config_test.go +++ b/deploy/snapshot/internal/types/config_test.go @@ -50,3 +50,12 @@ func TestAgentConfigValidateDefaultsStorageAccessMode(t *testing.T) { t.Fatalf("Storage.AccessMode = %q, want %q", cfg.Storage.AccessMode, StorageAccessModeAgentMount) } } + +func TestAgentConfigValidateRequiresPageBrokerControlSocket(t *testing.T) { + cfg := validAgentConfig() + cfg.PageBroker.Enabled = true + + if err := cfg.Validate(); err == nil { + t.Fatal("expected error for missing PageBroker control socket") + } +} diff --git a/deploy/snapshot/pagebroker/Makefile b/deploy/snapshot/pagebroker/Makefile index 4ea687d22933..0a1076fbce8f 100644 --- a/deploy/snapshot/pagebroker/Makefile +++ b/deploy/snapshot/pagebroker/Makefile @@ -1,7 +1,7 @@ PROTO := v1/pagebroker.proto -GTEST_FLAGS := $(shell pkg-config --cflags --libs gtest_main) -BROKER_SOURCES := broker.cpp checkpoint_transaction_descriptor.cpp posix_copy_engine.cpp restore_transaction_descriptor.cpp transfer_engine.cpp -DAEMON_SOURCES := $(BROKER_SOURCES) daemon.cpp file_descriptor.cpp +GTEST_FLAGS = $(shell pkg-config --cflags --libs gtest_main) +BROKER_SOURCES := broker.cpp checkpoint_transaction_descriptor.cpp posix_copy_engine.cpp restore_transaction_descriptor.cpp transaction.cpp transfer_engine.cpp +DAEMON_SOURCES := $(BROKER_SOURCES) connection_tracker.cpp daemon.cpp file_descriptor.cpp .PHONY: daemon generate test diff --git a/deploy/snapshot/pagebroker/broker.cpp b/deploy/snapshot/pagebroker/broker.cpp index 8c10c27388ee..2a2f2e9e5651 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/pagebroker/broker.cpp @@ -98,6 +98,30 @@ Broker::Broker(Path staging_root) : staging_root_(fs::weakly_canonical(std::move fs::create_directories(staging_root_ / "checkpoint"); } +Transaction& +Broker::GetTransaction(const std::string& transaction_id) +{ + std::lock_guard lock(transactions_mutex_); + return transactions_.try_emplace(transaction_id).first->second; +} + +bool +Broker::ReserveStaging(uintmax_t bytes) +{ + std::lock_guard lock(transactions_mutex_); + if (!HasAvailableSpace(staging_root_, bytes + reserved_staging_bytes_)) + return false; + reserved_staging_bytes_ += bytes; + return true; +} + +void +Broker::ReleaseStaging(uintmax_t bytes) +{ + std::lock_guard lock(transactions_mutex_); + reserved_staging_bytes_ -= bytes; +} + const TransferEngine& Broker::Engine(TransferEngineType engine_type) const { @@ -147,20 +171,28 @@ Broker::Restore(const Request& request) 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())) + Transaction& transaction = GetTransaction(request.transaction_id()); + std::lock_guard lock(transaction.mutex()); + if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) return Fail(request, Failure::TRANSACTION_CONFLICT, "restore transaction conflicts"); - if (!HasAvailableSpace(staging_root_, TreeSize(source))) + const uintmax_t bytes = TreeSize(source); + if (!ReserveStaging(bytes)) return Fail(request, Failure::INSUFFICIENT_STORAGE, "insufficient tmpfs capacity"); - try { - transaction_states_.emplace(request.transaction_id(), TransactionState::LIVE); + transaction.set_state(Transaction::State::PREPARING); + transaction.set_reserved_bytes(bytes); engine.CopyDirectory(source, staging_directory); - restore_transactions_.emplace(request.transaction_id(), RestoreTransactionDescriptor(staging_directory)); + ReleaseStaging(bytes); + transaction.set_reserved_bytes(0); + transaction.set_descriptor(RestoreTransactionDescriptor(staging_directory)); + transaction.set_state(Transaction::State::STAGED); } catch (...) { + ReleaseStaging(transaction.reserved_bytes()); + transaction.set_reserved_bytes(0); fs::remove_all(staging_directory); - transaction_states_.erase(request.transaction_id()); - restore_transactions_.erase(request.transaction_id()); + transaction.clear_descriptor(); + transaction.set_state(Transaction::State::NEW); throw; } auto response = Reply(request); @@ -181,19 +213,20 @@ Broker::PrepareCheckpoint(const Request& request) 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())) + Transaction& transaction = GetTransaction(request.transaction_id()); + std::lock_guard lock(transaction.mutex()); + if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) return Fail(request, Failure::TRANSACTION_CONFLICT, "checkpoint transaction conflicts"); try { - transaction_states_.emplace(request.transaction_id(), TransactionState::LIVE); + transaction.set_state(Transaction::State::PREPARING); fs::create_directory(staging_directory); - checkpoint_transactions_.emplace( - request.transaction_id(), - CheckpointTransactionDescriptor(staging_directory, operation.destination(), engine.type())); + transaction.set_descriptor(CheckpointTransactionDescriptor(staging_directory, operation.destination(), engine.type())); + transaction.set_state(Transaction::State::STAGED); } catch (...) { fs::remove_all(staging_directory); - transaction_states_.erase(request.transaction_id()); - checkpoint_transactions_.erase(request.transaction_id()); + transaction.clear_descriptor(); + transaction.set_state(Transaction::State::NEW); throw; } auto response = Reply(request); @@ -204,29 +237,32 @@ Broker::PrepareCheckpoint(const Request& request) Response Broker::Commit(const Request& request) { - const auto state = transaction_states_.find(request.transaction_id()); - if (state == transaction_states_.end() || state->second == TransactionState::ABORTED) + Transaction& transaction = GetTransaction(request.transaction_id()); + std::lock_guard lock(transaction.mutex()); + if (transaction.state() == Transaction::State::NEW || transaction.state() == Transaction::State::ABORTED) return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); - if (state->second == TransactionState::COMMITTED) + if (transaction.state() == Transaction::State::PREPARING) + return Fail(request, Failure::TRANSACTION_CONFLICT, "transaction is preparing"); + if (transaction.state() == Transaction::State::COMMITTED) return CommitSucceeded(request); - const auto restore = restore_transactions_.find(request.transaction_id()); - if (restore != restore_transactions_.end()) - return CleanupRestore(request, restore->second); + if (const auto* restore = std::get_if(&transaction.descriptor())) + return CleanupRestore(request, *restore); - const auto checkpoint = checkpoint_transactions_.find(request.transaction_id()); - if (checkpoint == checkpoint_transactions_.end()) + const auto* checkpoint = std::get_if(&transaction.descriptor()); + if (checkpoint == nullptr) return Fail(request, Failure::INTERNAL_ERROR, "live transaction has no descriptor"); - return PublishCheckpoint(request, checkpoint->second); + return PublishCheckpoint(request, *checkpoint); } 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; + Transaction& current = GetTransaction(request.transaction_id()); + current.clear_descriptor(); + current.set_state(Transaction::State::COMMITTED); return CommitSucceeded(request); } @@ -248,8 +284,9 @@ Broker::PublishCheckpoint(const Request& request, const CheckpointTransactionDes 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; + Transaction& current = GetTransaction(request.transaction_id()); + current.clear_descriptor(); + current.set_state(Transaction::State::COMMITTED); std::error_code cleanup_error; fs::remove_all(staging_directory, cleanup_error); } @@ -263,19 +300,19 @@ Broker::PublishCheckpoint(const Request& request, const CheckpointTransactionDes Response Broker::Abort(const Request& request) { - const auto state = transaction_states_.find(request.transaction_id()); - if (state == transaction_states_.end() || state->second == TransactionState::COMMITTED) + Transaction& transaction = GetTransaction(request.transaction_id()); + std::lock_guard lock(transaction.mutex()); + if (transaction.state() == Transaction::State::NEW || transaction.state() == Transaction::State::COMMITTED) return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); - if (state->second == TransactionState::ABORTED) + if (transaction.state() == Transaction::State::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; + transaction.clear_descriptor(); + transaction.set_state(Transaction::State::ABORTED); return AbortSucceeded(request); } diff --git a/deploy/snapshot/pagebroker/broker.hpp b/deploy/snapshot/pagebroker/broker.hpp index 2fde5f25229f..b0244079f265 100644 --- a/deploy/snapshot/pagebroker/broker.hpp +++ b/deploy/snapshot/pagebroker/broker.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -8,6 +9,7 @@ #include "pagebroker_types.hpp" #include "restore_transaction_descriptor.hpp" #include "transfer_engine.hpp" +#include "transaction.hpp" namespace snapshot::pagebroker { class Broker { @@ -17,12 +19,12 @@ class Broker { 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; + using Transactions = std::unordered_map; const TransferEngine& Engine(TransferEngineType engine_type) const; + Transaction& GetTransaction(const std::string& transaction_id); + bool ReserveStaging(uintmax_t bytes); + void ReleaseStaging(uintmax_t bytes); 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. @@ -32,8 +34,8 @@ class Broker { Response Abort(const Request& request); Path staging_root_; Engines io_engines_; - CheckpointTransactions checkpoint_transactions_; - RestoreTransactions restore_transactions_; - TransactionStates transaction_states_; + std::mutex transactions_mutex_; + Transactions transactions_; + uintmax_t reserved_staging_bytes_ = 0; }; } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/connection_tracker.cpp b/deploy/snapshot/pagebroker/connection_tracker.cpp new file mode 100644 index 000000000000..f3836eb7a05f --- /dev/null +++ b/deploy/snapshot/pagebroker/connection_tracker.cpp @@ -0,0 +1,23 @@ +#include "connection_tracker.hpp" + +void +ConnectionTracker::Start() +{ + std::lock_guard lock(mutex_); + ++active_; +} + +void +ConnectionTracker::Finish() +{ + std::lock_guard lock(mutex_); + if (--active_ == 0) + finished_.notify_all(); +} + +void +ConnectionTracker::Wait() +{ + std::unique_lock lock(mutex_); + finished_.wait(lock, [this] { return active_ == 0; }); +} diff --git a/deploy/snapshot/pagebroker/connection_tracker.hpp b/deploy/snapshot/pagebroker/connection_tracker.hpp new file mode 100644 index 000000000000..10d3ba1c079b --- /dev/null +++ b/deploy/snapshot/pagebroker/connection_tracker.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +class ConnectionTracker { + public: + void Start(); + void Finish(); + void Wait(); + + private: + std::mutex mutex_; + std::condition_variable finished_; + unsigned int active_ = 0; +}; diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index 4131aed609ea..ebcf6f88731e 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -1,5 +1,7 @@ #include #include +#include +#include #include #include #include @@ -8,8 +10,11 @@ #include #include #include +#include +#include #include "broker.hpp" +#include "connection_tracker.hpp" #include "file_descriptor.hpp" namespace fs = std::filesystem; @@ -20,6 +25,8 @@ using snapshot::pagebroker::Response; namespace { constexpr uint32_t kMaxMessageSize = 64 << 10; // 64 KB +constexpr timeval kConnectionTimeout{30, 0}; +constexpr int kShutdownPollTimeoutMs = 1000; enum ArgumentIndex { kSocketPath = 1, kStagingDirectory, kArgumentCount }; volatile sig_atomic_t shutting_down; @@ -100,6 +107,22 @@ HandleConnection(int connection, Broker& broker) WriteAll(connection, &size, sizeof(size)); WriteAll(connection, message.data(), message.size()); } + +void +ServeConnection(int connection, Broker& broker, ConnectionTracker& connections) +{ + FileDescriptor descriptor(connection); + try { + if (setsockopt(descriptor.get(), SOL_SOCKET, SO_RCVTIMEO, &kConnectionTimeout, sizeof(kConnectionTimeout)) < 0 || + setsockopt(descriptor.get(), SOL_SOCKET, SO_SNDTIMEO, &kConnectionTimeout, sizeof(kConnectionTimeout)) < 0) + throw std::system_error(errno, std::generic_category(), "set connection timeout"); + HandleConnection(descriptor.get(), broker); + } + catch (const std::exception& error) { + std::cerr << "handle connection: " << error.what() << '\n'; + } + connections.Finish(); +} } // namespace int @@ -149,15 +172,33 @@ main(int argc, char** argv) std::cerr << "listen: " << std::strerror(errno) << '\n'; return 1; } + const int listener_flags = fcntl(listener.get(), F_GETFL); + if (listener_flags < 0 || fcntl(listener.get(), F_SETFL, listener_flags | O_NONBLOCK) < 0) { + std::cerr << "make listener nonblocking: " << std::strerror(errno) << '\n'; + return 1; + } Broker broker(argv[kStagingDirectory]); + ConnectionTracker connections; while (!shutting_down) { - FileDescriptor connection(accept(listener.get(), nullptr, nullptr)); - if (connection.get() < 0) { + pollfd poll_descriptor{listener.get(), POLLIN, 0}; + if (poll(&poll_descriptor, 1, kShutdownPollTimeoutMs) <= 0) + continue; + const int connection = accept(listener.get(), nullptr, nullptr); + if (connection < 0) { if (errno != EINTR) std::cerr << "accept: " << std::strerror(errno) << '\n'; continue; } - HandleConnection(connection.get(), broker); + connections.Start(); + try { + std::thread(ServeConnection, connection, std::ref(broker), std::ref(connections)).detach(); + } + catch (const std::system_error& error) { + FileDescriptor descriptor(connection); + connections.Finish(); + std::cerr << "start connection: " << error.what() << '\n'; + } } + connections.Wait(); } diff --git a/deploy/snapshot/pagebroker/daemon_test.cpp b/deploy/snapshot/pagebroker/daemon_test.cpp index a6da7ca2c7af..b795f3b1b3a6 100644 --- a/deploy/snapshot/pagebroker/daemon_test.cpp +++ b/deploy/snapshot/pagebroker/daemon_test.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "broker.hpp" @@ -77,6 +78,25 @@ TEST_F(BrokerTest, StagesRestoreAndCleansUpOnCommit) EXPECT_EQ(abort_response.failure().code(), Failure::TRANSACTION_NOT_FOUND); } +TEST_F(BrokerTest, StagesIndependentRestoresConcurrently) +{ + auto first = RequestFor("first"); + auto second = RequestFor("second"); + Configure(first.mutable_staged_restore()->mutable_source(), first.mutable_staged_restore()->mutable_io_engine(), source_); + Configure(second.mutable_staged_restore()->mutable_source(), second.mutable_staged_restore()->mutable_io_engine(), source_); + + Response first_response; + Response second_response; + std::thread first_request([&] { first_response = broker().HandleRequest(first); }); + std::thread second_request([&] { second_response = broker().HandleRequest(second); }); + first_request.join(); + second_request.join(); + + ASSERT_TRUE(first_response.has_staged_restore_directory()); + ASSERT_TRUE(second_response.has_staged_restore_directory()); + EXPECT_NE(first_response.staged_restore_directory().image_directory(), second_response.staged_restore_directory().image_directory()); +} + TEST_F(BrokerTest, RejectsUnsafeTransactionIDs) { for (const auto& id : diff --git a/deploy/snapshot/pagebroker/transaction.cpp b/deploy/snapshot/pagebroker/transaction.cpp new file mode 100644 index 000000000000..9b0d5a937a2c --- /dev/null +++ b/deploy/snapshot/pagebroker/transaction.cpp @@ -0,0 +1,53 @@ +#include "transaction.hpp" + +#include + +namespace snapshot::pagebroker { +std::mutex& +Transaction::mutex() +{ + return mutex_; +} + +Transaction::State +Transaction::state() const +{ + return state_; +} + +void +Transaction::set_state(State state) +{ + state_ = state; +} + +const Transaction::Descriptor& +Transaction::descriptor() const +{ + return descriptor_; +} + +void +Transaction::set_descriptor(Descriptor descriptor) +{ + descriptor_ = std::move(descriptor); +} + +void +Transaction::clear_descriptor() +{ + descriptor_ = std::monostate(); +} + +uintmax_t +Transaction::reserved_bytes() const +{ + return reserved_bytes_; +} + +void +Transaction::set_reserved_bytes(uintmax_t bytes) +{ + reserved_bytes_ = bytes; +} +} // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/transaction.hpp b/deploy/snapshot/pagebroker/transaction.hpp new file mode 100644 index 000000000000..7cf75bd0b617 --- /dev/null +++ b/deploy/snapshot/pagebroker/transaction.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include +#include +#include + +#include "checkpoint_transaction_descriptor.hpp" +#include "restore_transaction_descriptor.hpp" + +namespace snapshot::pagebroker { +class Transaction { + public: + enum class State { NEW, PREPARING, STAGED, COMMITTED, ABORTED }; + using Descriptor = std::variant; + + std::mutex& mutex(); + State state() const; + void set_state(State state); + const Descriptor& descriptor() const; + void set_descriptor(Descriptor descriptor); + void clear_descriptor(); + uintmax_t reserved_bytes() const; + void set_reserved_bytes(uintmax_t bytes); + + private: + std::mutex mutex_; + State state_ = State::NEW; + Descriptor descriptor_; + uintmax_t reserved_bytes_ = 0; +}; +} // namespace snapshot::pagebroker From e6b4d0ca3f366b230168ab92bee9ecabb51b4bc9 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 14:40:00 +0300 Subject: [PATCH 02/10] feat(snapshot): complete concurrent PageBroker staging --- deploy/snapshot/pagebroker/Makefile | 2 +- deploy/snapshot/pagebroker/broker.cpp | 123 ++++---- deploy/snapshot/pagebroker/broker.hpp | 8 +- .../pagebroker/connection_tracker.cpp | 23 -- .../pagebroker/connection_tracker.hpp | 16 -- deploy/snapshot/pagebroker/daemon.cpp | 262 +++++++++++++----- deploy/snapshot/pagebroker/daemon.hpp | 7 + deploy/snapshot/pagebroker/daemon_test.cpp | 60 +++- .../snapshot/pagebroker/file_descriptor.cpp | 15 + .../snapshot/pagebroker/file_descriptor.hpp | 2 + deploy/snapshot/pagebroker/main.cpp | 13 + .../snapshot/pagebroker/posix_copy_engine.cpp | 82 ++++++ .../snapshot/pagebroker/posix_copy_engine.hpp | 4 + deploy/snapshot/pagebroker/transaction.cpp | 11 - deploy/snapshot/pagebroker/transaction.hpp | 4 - .../snapshot/pagebroker/transfer_engine.hpp | 7 + 16 files changed, 456 insertions(+), 183 deletions(-) delete mode 100644 deploy/snapshot/pagebroker/connection_tracker.cpp delete mode 100644 deploy/snapshot/pagebroker/connection_tracker.hpp create mode 100644 deploy/snapshot/pagebroker/daemon.hpp create mode 100644 deploy/snapshot/pagebroker/main.cpp diff --git a/deploy/snapshot/pagebroker/Makefile b/deploy/snapshot/pagebroker/Makefile index 0a1076fbce8f..524551e0ba01 100644 --- a/deploy/snapshot/pagebroker/Makefile +++ b/deploy/snapshot/pagebroker/Makefile @@ -1,7 +1,7 @@ PROTO := v1/pagebroker.proto GTEST_FLAGS = $(shell pkg-config --cflags --libs gtest_main) BROKER_SOURCES := broker.cpp checkpoint_transaction_descriptor.cpp posix_copy_engine.cpp restore_transaction_descriptor.cpp transaction.cpp transfer_engine.cpp -DAEMON_SOURCES := $(BROKER_SOURCES) connection_tracker.cpp daemon.cpp file_descriptor.cpp +DAEMON_SOURCES := $(BROKER_SOURCES) daemon.cpp main.cpp file_descriptor.cpp .PHONY: daemon generate test diff --git a/deploy/snapshot/pagebroker/broker.cpp b/deploy/snapshot/pagebroker/broker.cpp index 2a2f2e9e5651..366fdf9f6414 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/pagebroker/broker.cpp @@ -55,23 +55,29 @@ IsSafePathComponent(const std::string& value) value.find('\\') == std::string::npos && value.find('\0') == std::string::npos; } -bool -IsFilesystemPosix(const StorageBackend& storage, const IOEngine& engine) +const StorageBackend& +ValidateStagedRestore(const StagedRestoreRequest& request) { - return storage.has_filesystem() && !storage.filesystem().directory().empty() && engine.has_posix_copy(); + if (!request.has_source() || request.source().kind_case() == StorageBackend::KIND_NOT_SET) + throw std::invalid_argument("restore source is required"); + return request.source(); } -uintmax_t -TreeSize(const Path& path) +const StorageBackend& +ValidateStagedCheckpoint(const PrepareStagedCheckpointRequest& request) { - uintmax_t bytes = 0; - for (const auto& entry : fs::recursive_directory_iterator(path)) { + if (!request.has_destination() || request.destination().kind_case() == StorageBackend::KIND_NOT_SET) + throw std::invalid_argument("checkpoint destination is required"); + return request.destination(); +} + +void +RejectSymlinks(const Path& directory) +{ + for (const auto& entry : fs::recursive_directory_iterator(directory)) { if (entry.is_symlink()) throw std::runtime_error("checkpoint contains symlink"); - if (entry.is_regular_file()) - bytes += entry.file_size(); } - return bytes; } bool @@ -122,6 +128,19 @@ Broker::ReleaseStaging(uintmax_t bytes) reserved_staging_bytes_ -= bytes; } +Response +Broker::AbortStaging( + const Request& request, Transaction& transaction, const Path& staging_directory, const std::exception& error) +{ + std::error_code cleanup_error; + fs::remove_all(staging_directory, cleanup_error); + if (cleanup_error) + return Fail(request, Failure::STORAGE_ERROR, std::string(error.what()) + "; cleanup: " + cleanup_error.message()); + transaction.clear_descriptor(); + transaction.set_state(Transaction::State::ABORTED); + return Fail(request, Failure::STORAGE_ERROR, error.what()); +} + const TransferEngine& Broker::Engine(TransferEngineType engine_type) const { @@ -132,6 +151,14 @@ Broker::Engine(TransferEngineType engine_type) const throw std::runtime_error("configured I/O engine not found"); } +const TransferEngine& +Broker::Engine(const IOEngine& engine) const +{ + if (engine.has_posix_copy()) + return Engine(TransferEngineType::POSIX_COPY); + throw std::invalid_argument("unsupported I/O engine"); +} + Response Broker::HandleRequest(const Request& request) { @@ -153,6 +180,9 @@ Broker::HandleRequest(const Request& request) return Fail(request, Failure::INVALID_REQUEST, "unsupported operation"); } } + catch (const std::invalid_argument& error) { + return Fail(request, Failure::INVALID_REQUEST, error.what()); + } catch (const std::exception& error) { return Fail(request, Failure::STORAGE_ERROR, error.what()); } @@ -162,38 +192,36 @@ 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 auto& source = ValidateStagedRestore(operation); + const auto& engine = Engine(operation.io_engine()); + return StageRestore(request, source, engine); +} - const Path source(operation.source().filesystem().directory()); +Response +Broker::StageRestore(const Request& request, const StorageBackend& source, const TransferEngine& engine) +{ 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"); + const uintmax_t bytes = engine.RestoreSize(source); Transaction& transaction = GetTransaction(request.transaction_id()); std::lock_guard lock(transaction.mutex()); if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) return Fail(request, Failure::TRANSACTION_CONFLICT, "restore transaction conflicts"); - const uintmax_t bytes = TreeSize(source); if (!ReserveStaging(bytes)) return Fail(request, Failure::INSUFFICIENT_STORAGE, "insufficient tmpfs capacity"); + bool staging_reserved = true; try { transaction.set_state(Transaction::State::PREPARING); - transaction.set_reserved_bytes(bytes); - engine.CopyDirectory(source, staging_directory); + engine.StageRestore(source, staging_directory); ReleaseStaging(bytes); - transaction.set_reserved_bytes(0); + staging_reserved = false; transaction.set_descriptor(RestoreTransactionDescriptor(staging_directory)); transaction.set_state(Transaction::State::STAGED); } - catch (...) { - ReleaseStaging(transaction.reserved_bytes()); - transaction.set_reserved_bytes(0); - fs::remove_all(staging_directory); - transaction.clear_descriptor(); - transaction.set_state(Transaction::State::NEW); - throw; + catch (const std::exception& error) { + if (staging_reserved) + ReleaseStaging(bytes); + return AbortStaging(request, transaction, staging_directory, error); } auto response = Reply(request); response.mutable_staged_restore_directory()->set_image_directory(staging_directory.string()); @@ -204,15 +232,16 @@ 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 auto& destination = ValidateStagedCheckpoint(operation); + const auto& engine = Engine(operation.io_engine()); + return StageCheckpoint(request, destination, engine); +} - const Path destination(operation.destination().filesystem().directory()); +Response +Broker::StageCheckpoint(const Request& request, const StorageBackend& destination, const TransferEngine& engine) +{ 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"); Transaction& transaction = GetTransaction(request.transaction_id()); std::lock_guard lock(transaction.mutex()); if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) @@ -220,14 +249,11 @@ Broker::PrepareCheckpoint(const Request& request) try { transaction.set_state(Transaction::State::PREPARING); fs::create_directory(staging_directory); - transaction.set_descriptor(CheckpointTransactionDescriptor(staging_directory, operation.destination(), engine.type())); + transaction.set_descriptor(CheckpointTransactionDescriptor(staging_directory, destination, engine.type())); transaction.set_state(Transaction::State::STAGED); } - catch (...) { - fs::remove_all(staging_directory); - transaction.clear_descriptor(); - transaction.set_state(Transaction::State::NEW); - throw; + catch (const std::exception& error) { + return AbortStaging(request, transaction, staging_directory, error); } auto response = Reply(request); response.mutable_staged_checkpoint_directory()->set_image_directory(staging_directory.string()); @@ -272,27 +298,20 @@ Broker::PublishCheckpoint(const Request& request, const CheckpointTransactionDes 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)) + const auto& engine = Engine(transaction.engine_type()); + if (engine.CheckpointDestinationConflicts(transaction.destination_storage())) return Fail(request, Failure::TRANSACTION_CONFLICT, "checkpoint destination conflicts"); - TreeSize(staging_directory); - + RejectSymlinks(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); + engine.PublishCheckpoint(staging_directory, transaction.destination_storage()); Transaction& current = GetTransaction(request.transaction_id()); current.clear_descriptor(); current.set_state(Transaction::State::COMMITTED); std::error_code cleanup_error; fs::remove_all(staging_directory, cleanup_error); } - catch (...) { - fs::remove_all(partial); - throw; + catch (const std::exception& error) { + return Fail(request, Failure::STORAGE_ERROR, error.what()); } return CommitSucceeded(request); } diff --git a/deploy/snapshot/pagebroker/broker.hpp b/deploy/snapshot/pagebroker/broker.hpp index b0244079f265..16c64fd85390 100644 --- a/deploy/snapshot/pagebroker/broker.hpp +++ b/deploy/snapshot/pagebroker/broker.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -8,8 +9,8 @@ #include "checkpoint_transaction_descriptor.hpp" #include "pagebroker_types.hpp" #include "restore_transaction_descriptor.hpp" -#include "transfer_engine.hpp" #include "transaction.hpp" +#include "transfer_engine.hpp" namespace snapshot::pagebroker { class Broker { @@ -22,11 +23,16 @@ class Broker { using Transactions = std::unordered_map; const TransferEngine& Engine(TransferEngineType engine_type) const; + const TransferEngine& Engine(const IOEngine& engine) const; Transaction& GetTransaction(const std::string& transaction_id); bool ReserveStaging(uintmax_t bytes); void ReleaseStaging(uintmax_t bytes); + Response AbortStaging( + const Request& request, Transaction& transaction, const Path& staging_directory, const std::exception& error); Response Restore(const Request& request); + Response StageRestore(const Request& request, const StorageBackend& source, const TransferEngine& engine); Response PrepareCheckpoint(const Request& request); + Response StageCheckpoint(const Request& request, const StorageBackend& destination, const TransferEngine& engine); // 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); diff --git a/deploy/snapshot/pagebroker/connection_tracker.cpp b/deploy/snapshot/pagebroker/connection_tracker.cpp deleted file mode 100644 index f3836eb7a05f..000000000000 --- a/deploy/snapshot/pagebroker/connection_tracker.cpp +++ /dev/null @@ -1,23 +0,0 @@ -#include "connection_tracker.hpp" - -void -ConnectionTracker::Start() -{ - std::lock_guard lock(mutex_); - ++active_; -} - -void -ConnectionTracker::Finish() -{ - std::lock_guard lock(mutex_); - if (--active_ == 0) - finished_.notify_all(); -} - -void -ConnectionTracker::Wait() -{ - std::unique_lock lock(mutex_); - finished_.wait(lock, [this] { return active_ == 0; }); -} diff --git a/deploy/snapshot/pagebroker/connection_tracker.hpp b/deploy/snapshot/pagebroker/connection_tracker.hpp deleted file mode 100644 index 10d3ba1c079b..000000000000 --- a/deploy/snapshot/pagebroker/connection_tracker.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include -#include - -class ConnectionTracker { - public: - void Start(); - void Finish(); - void Wait(); - - private: - std::mutex mutex_; - std::condition_variable finished_; - unsigned int active_ = 0; -}; diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index ebcf6f88731e..bfeadb006603 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -1,3 +1,5 @@ +#include "daemon.hpp" + #include #include #include @@ -6,15 +8,19 @@ #include #include +#include #include #include +#include #include #include +#include +#include #include -#include +#include +#include #include "broker.hpp" -#include "connection_tracker.hpp" #include "file_descriptor.hpp" namespace fs = std::filesystem; @@ -27,7 +33,6 @@ namespace { constexpr uint32_t kMaxMessageSize = 64 << 10; // 64 KB constexpr timeval kConnectionTimeout{30, 0}; constexpr int kShutdownPollTimeoutMs = 1000; -enum ArgumentIndex { kSocketPath = 1, kStagingDirectory, kArgumentCount }; volatile sig_atomic_t shutting_down; void @@ -36,6 +41,89 @@ Stop(int) shutting_down = 1; } +void +LogError(std::string_view operation, const std::error_code& error) +{ + std::cerr << operation << ": " << error.message() << '\n'; +} + +ExitCode +Fail(std::string_view operation, const std::error_code& error) +{ + LogError(operation, error); + return ExitCode::FAILURE; +} + +std::error_code +InstallSignalHandler(int signal) +{ + struct sigaction action {}; + action.sa_handler = Stop; + sigemptyset(&action.sa_mask); + if (sigaction(signal, &action, nullptr) < 0) + return {errno, std::generic_category()}; + return {}; +} + +std::error_code +InstallSignalHandlers() +{ + if (const auto error = InstallSignalHandler(SIGINT); error) + return error; + return InstallSignalHandler(SIGTERM); +} + +std::error_code +PrepareDirectories(const fs::path& socket_path, const fs::path& staging_directory) +{ + std::error_code error; + fs::create_directories(socket_path.parent_path(), error); + if (error) + return error; + fs::create_directories(staging_directory, error); + return error; +} + +std::error_code +ConfigureConnection(int connection) +{ + if (setsockopt(connection, SOL_SOCKET, SO_RCVTIMEO, &kConnectionTimeout, sizeof(kConnectionTimeout)) < 0 || + setsockopt(connection, SOL_SOCKET, SO_SNDTIMEO, &kConnectionTimeout, sizeof(kConnectionTimeout)) < 0) + return {errno, std::generic_category()}; + return {}; +} + +std::error_code +ConfigureListener(int listener) +{ + const int flags = fcntl(listener, F_GETFL); + if (flags < 0 || fcntl(listener, F_SETFL, flags | O_NONBLOCK) < 0) + return {errno, std::generic_category()}; + return {}; +} + +std::pair +CreateListener(const fs::path& socket_path) +{ + std::error_code error; + fs::remove(socket_path, error); + if (error) + return std::make_pair(FileDescriptor(-1), error); + + FileDescriptor listener(socket(AF_UNIX, SOCK_STREAM, 0)); + if (listener.get() < 0) + return std::make_pair(std::move(listener), std::error_code(errno, std::generic_category())); + sockaddr_un address{}; + address.sun_family = AF_UNIX; + std::strcpy(address.sun_path, socket_path.c_str()); + if (bind(listener.get(), reinterpret_cast(&address), sizeof(address)) < 0 || + listen(listener.get(), SOMAXCONN) < 0) + return std::make_pair(std::move(listener), std::error_code(errno, std::generic_category())); + if (error = ConfigureListener(listener.get()); error) + return std::make_pair(std::move(listener), error); + return std::make_pair(std::move(listener), std::error_code{}); +} + bool ReadAll(int fd, void* buffer, size_t size) { @@ -81,6 +169,42 @@ InvalidRequest() return response; } +const char* +CommandName(Request::CommandCase command) +{ + switch (command) { + case Request::kStagedRestore: + return "staged_restore"; + case Request::kPrepareStagedCheckpoint: + return "prepare_staged_checkpoint"; + case Request::kCommit: + return "commit"; + case Request::kAbort: + return "abort"; + default: + return "invalid"; + } +} + +const char* +ResultName(const Response& response) +{ + switch (response.result_case()) { + case Response::kStagedRestoreDirectory: + return "staged_restore"; + case Response::kStagedCheckpointDirectory: + return "staged_checkpoint"; + case Response::kCommitComplete: + return "committed"; + case Response::kAbortComplete: + return "aborted"; + case Response::kFailure: + return "failed"; + default: + return "invalid"; + } +} + void HandleConnection(int connection, Broker& broker) { @@ -99,6 +223,10 @@ HandleConnection(int connection, Broker& broker) response = InvalidRequest(); } else { response = broker.HandleRequest(request); + std::osyncstream(std::cerr) << "transaction=" << request.transaction_id() + << " command=" << CommandName(request.command_case()) + << " result=" << ResultName(response) + << (response.has_failure() ? " error=" + response.failure().message() : "") << '\n'; } } @@ -109,96 +237,86 @@ HandleConnection(int connection, Broker& broker) } void -ServeConnection(int connection, Broker& broker, ConnectionTracker& connections) +ServeConnection(int connection, Broker& broker) { FileDescriptor descriptor(connection); - try { - if (setsockopt(descriptor.get(), SOL_SOCKET, SO_RCVTIMEO, &kConnectionTimeout, sizeof(kConnectionTimeout)) < 0 || - setsockopt(descriptor.get(), SOL_SOCKET, SO_SNDTIMEO, &kConnectionTimeout, sizeof(kConnectionTimeout)) < 0) - throw std::system_error(errno, std::generic_category(), "set connection timeout"); - HandleConnection(descriptor.get(), broker); - } - catch (const std::exception& error) { - std::cerr << "handle connection: " << error.what() << '\n'; + if (const auto error = ConfigureConnection(descriptor.get()); error) { + LogError("set connection timeout", error); + return; } - connections.Finish(); + HandleConnection(descriptor.get(), broker); } -} // namespace -int -main(int argc, char** argv) +void +ReapHandlers(std::vector>& handlers) { - if (argc != kArgumentCount) { - std::cerr << "usage: pagebroker-daemon SOCKET STAGING_DIRECTORY\n"; - return 2; - } - - struct sigaction action {}; - action.sa_handler = Stop; - sigemptyset(&action.sa_mask); - if (sigaction(SIGINT, &action, nullptr) < 0 || sigaction(SIGTERM, &action, nullptr) < 0) { - std::cerr << "install signal handler: " << std::strerror(errno) << '\n'; - return 1; - } - - const fs::path socket_path(argv[kSocketPath]); - std::error_code error; - fs::create_directories(socket_path.parent_path(), error); - if (error) { - std::cerr << "create socket directory: " << error.message() << '\n'; - return 1; - } - fs::create_directories(argv[kStagingDirectory], error); - if (error) { - std::cerr << "create staging directory: " << error.message() << '\n'; - return 1; - } - if (socket_path.string().size() >= sizeof(sockaddr_un::sun_path)) { - std::cerr << "socket path is too long\n"; - return 2; + for (auto handler = handlers.begin(); handler != handlers.end();) { + if (handler->wait_for(std::chrono::seconds(0)) != std::future_status::ready) { + ++handler; + continue; + } + handler->get(); + handler = handlers.erase(handler); } - unlink(socket_path.c_str()); +} - FileDescriptor listener(socket(AF_UNIX, SOCK_STREAM, 0)); - if (listener.get() < 0) { - std::cerr << "create listener: " << std::strerror(errno) << '\n'; - return 1; - } - sockaddr_un address{}; - address.sun_family = AF_UNIX; - std::strcpy(address.sun_path, socket_path.c_str()); - if (bind(listener.get(), reinterpret_cast(&address), sizeof(address)) < 0 || - listen(listener.get(), 16) < 0) { - std::cerr << "listen: " << std::strerror(errno) << '\n'; - return 1; - } - const int listener_flags = fcntl(listener.get(), F_GETFL); - if (listener_flags < 0 || fcntl(listener.get(), F_SETFL, listener_flags | O_NONBLOCK) < 0) { - std::cerr << "make listener nonblocking: " << std::strerror(errno) << '\n'; - return 1; - } +void +WaitForHandlers(std::vector>& handlers) +{ + for (auto& handler : handlers) handler.get(); +} - Broker broker(argv[kStagingDirectory]); - ConnectionTracker connections; +void +Serve(FileDescriptor& listener, Broker& broker) +{ + std::vector> handlers; while (!shutting_down) { + ReapHandlers(handlers); pollfd poll_descriptor{listener.get(), POLLIN, 0}; - if (poll(&poll_descriptor, 1, kShutdownPollTimeoutMs) <= 0) + const int ready = poll(&poll_descriptor, 1, kShutdownPollTimeoutMs); + if (ready == 0) + continue; + if (ready < 0) { + if (errno != EINTR) + LogError("poll", {errno, std::generic_category()}); continue; + } const int connection = accept(listener.get(), nullptr, nullptr); if (connection < 0) { - if (errno != EINTR) - std::cerr << "accept: " << std::strerror(errno) << '\n'; + if (errno != EINTR && errno != EAGAIN && errno != EWOULDBLOCK) + LogError("accept", {errno, std::generic_category()}); continue; } - connections.Start(); try { - std::thread(ServeConnection, connection, std::ref(broker), std::ref(connections)).detach(); + handlers.emplace_back( + std::async(std::launch::async, [connection, &broker] { ServeConnection(connection, broker); })); } catch (const std::system_error& error) { FileDescriptor descriptor(connection); - connections.Finish(); std::cerr << "start connection: " << error.what() << '\n'; } } - connections.Wait(); + WaitForHandlers(handlers); +} +} // namespace + +ExitCode +RunDaemon(const fs::path& socket_path, const fs::path& staging_directory) +{ + shutting_down = 0; + if (const auto error = InstallSignalHandlers(); error) + return Fail("install signal handlers", error); + if (const auto error = PrepareDirectories(socket_path, staging_directory); error) + return Fail("create daemon directories", error); + if (socket_path.string().size() >= sizeof(sockaddr_un::sun_path)) { + std::cerr << "socket path is too long\n"; + return ExitCode::INVALID_ARGUMENTS; + } + auto [listener, error] = CreateListener(socket_path); + if (error) + return Fail("create listener", error); + + Broker broker(staging_directory); + Serve(listener, broker); + return ExitCode::SUCCESS; } diff --git a/deploy/snapshot/pagebroker/daemon.hpp b/deploy/snapshot/pagebroker/daemon.hpp new file mode 100644 index 000000000000..e9974ab9392b --- /dev/null +++ b/deploy/snapshot/pagebroker/daemon.hpp @@ -0,0 +1,7 @@ +#pragma once + +#include + +enum class ExitCode { SUCCESS = 0, FAILURE = 1, INVALID_ARGUMENTS = 2 }; + +ExitCode RunDaemon(const std::filesystem::path& socket_path, const std::filesystem::path& staging_directory); diff --git a/deploy/snapshot/pagebroker/daemon_test.cpp b/deploy/snapshot/pagebroker/daemon_test.cpp index b795f3b1b3a6..338209beb139 100644 --- a/deploy/snapshot/pagebroker/daemon_test.cpp +++ b/deploy/snapshot/pagebroker/daemon_test.cpp @@ -82,8 +82,10 @@ TEST_F(BrokerTest, StagesIndependentRestoresConcurrently) { auto first = RequestFor("first"); auto second = RequestFor("second"); - Configure(first.mutable_staged_restore()->mutable_source(), first.mutable_staged_restore()->mutable_io_engine(), source_); - Configure(second.mutable_staged_restore()->mutable_source(), second.mutable_staged_restore()->mutable_io_engine(), source_); + Configure( + first.mutable_staged_restore()->mutable_source(), first.mutable_staged_restore()->mutable_io_engine(), source_); + Configure( + second.mutable_staged_restore()->mutable_source(), second.mutable_staged_restore()->mutable_io_engine(), source_); Response first_response; Response second_response; @@ -94,7 +96,28 @@ TEST_F(BrokerTest, StagesIndependentRestoresConcurrently) ASSERT_TRUE(first_response.has_staged_restore_directory()); ASSERT_TRUE(second_response.has_staged_restore_directory()); - EXPECT_NE(first_response.staged_restore_directory().image_directory(), second_response.staged_restore_directory().image_directory()); + EXPECT_NE( + first_response.staged_restore_directory().image_directory(), + second_response.staged_restore_directory().image_directory()); +} + +TEST_F(BrokerTest, RejectsConcurrentRestoreForSameTransaction) +{ + auto first = RequestFor("restore"); + auto second = RequestFor("restore"); + Configure( + first.mutable_staged_restore()->mutable_source(), first.mutable_staged_restore()->mutable_io_engine(), source_); + Configure( + second.mutable_staged_restore()->mutable_source(), second.mutable_staged_restore()->mutable_io_engine(), source_); + + Response first_response; + Response second_response; + std::thread first_request([&] { first_response = broker().HandleRequest(first); }); + std::thread second_request([&] { second_response = broker().HandleRequest(second); }); + first_request.join(); + second_request.join(); + + EXPECT_NE(first_response.has_staged_restore_directory(), second_response.has_staged_restore_directory()); } TEST_F(BrokerTest, RejectsUnsafeTransactionIDs) @@ -122,6 +145,37 @@ TEST_F(BrokerTest, RejectsSymlinkInRestoreSource) EXPECT_EQ(response.failure().code(), Failure::STORAGE_ERROR); } +TEST_F(BrokerTest, RejectsInvalidStagedRestore) +{ + auto restore = RequestFor("restore"); + restore.mutable_staged_restore(); + const auto response = broker().HandleRequest(restore); + ASSERT_TRUE(response.has_failure()); + EXPECT_EQ(response.failure().code(), Failure::INVALID_REQUEST); +} + +TEST_F(BrokerTest, AbortsFailedCheckpointStaging) +{ + const fs::path destination = root_ / "storage" / "published"; + const fs::path staging_directory = root_ / "tmpfs" / "checkpoint" / "checkpoint"; + fs::create_symlink(root_ / "missing", staging_directory); + auto prepare = RequestFor("checkpoint"); + Configure( + prepare.mutable_prepare_staged_checkpoint()->mutable_destination(), + prepare.mutable_prepare_staged_checkpoint()->mutable_io_engine(), destination); + const auto failed = broker().HandleRequest(prepare); + ASSERT_TRUE(failed.has_failure()); + EXPECT_EQ(failed.failure().code(), Failure::STORAGE_ERROR); + + auto abort = RequestFor("checkpoint"); + abort.mutable_abort(); + EXPECT_TRUE(broker().HandleRequest(abort).has_abort_complete()); + + const auto retry = broker().HandleRequest(prepare); + ASSERT_TRUE(retry.has_failure()); + EXPECT_EQ(retry.failure().code(), Failure::TRANSACTION_CONFLICT); +} + TEST_F(BrokerTest, PublishesCheckpoint) { const fs::path published = root_ / "storage" / "published"; diff --git a/deploy/snapshot/pagebroker/file_descriptor.cpp b/deploy/snapshot/pagebroker/file_descriptor.cpp index 88993926c7f3..074adae2de3c 100644 --- a/deploy/snapshot/pagebroker/file_descriptor.cpp +++ b/deploy/snapshot/pagebroker/file_descriptor.cpp @@ -2,6 +2,8 @@ #include +#include + FileDescriptor::FileDescriptor(int value) : value_(value) {} FileDescriptor::~FileDescriptor() noexcept @@ -10,6 +12,19 @@ FileDescriptor::~FileDescriptor() noexcept close(value_); } +FileDescriptor::FileDescriptor(FileDescriptor&& other) noexcept : value_(std::exchange(other.value_, -1)) {} + +FileDescriptor& +FileDescriptor::operator=(FileDescriptor&& other) noexcept +{ + if (this != &other) { + if (value_ >= 0) + close(value_); + value_ = std::exchange(other.value_, -1); + } + return *this; +} + int FileDescriptor::get() const { diff --git a/deploy/snapshot/pagebroker/file_descriptor.hpp b/deploy/snapshot/pagebroker/file_descriptor.hpp index b178cf540739..c858259a6129 100644 --- a/deploy/snapshot/pagebroker/file_descriptor.hpp +++ b/deploy/snapshot/pagebroker/file_descriptor.hpp @@ -7,6 +7,8 @@ class FileDescriptor { FileDescriptor(const FileDescriptor&) = delete; FileDescriptor& operator=(const FileDescriptor&) = delete; + FileDescriptor(FileDescriptor&& other) noexcept; + FileDescriptor& operator=(FileDescriptor&& other) noexcept; int get() const; diff --git a/deploy/snapshot/pagebroker/main.cpp b/deploy/snapshot/pagebroker/main.cpp new file mode 100644 index 000000000000..22f274f5967a --- /dev/null +++ b/deploy/snapshot/pagebroker/main.cpp @@ -0,0 +1,13 @@ +#include + +#include "daemon.hpp" + +int +main(int argc, char** argv) +{ + if (argc != 3) { + std::cerr << "usage: pagebroker-daemon SOCKET STAGING_DIRECTORY\n"; + return static_cast(ExitCode::INVALID_ARGUMENTS); + } + return static_cast(RunDaemon(argv[1], argv[2])); +} diff --git a/deploy/snapshot/pagebroker/posix_copy_engine.cpp b/deploy/snapshot/pagebroker/posix_copy_engine.cpp index 44be57509ba1..91dd2114d94e 100644 --- a/deploy/snapshot/pagebroker/posix_copy_engine.cpp +++ b/deploy/snapshot/pagebroker/posix_copy_engine.cpp @@ -1,14 +1,96 @@ #include "posix_copy_engine.hpp" #include +#include namespace snapshot::pagebroker { +namespace { +Path +SourcePath(const StorageBackend& source) +{ + if (!source.has_filesystem() || source.filesystem().directory().empty()) + throw std::invalid_argument("filesystem source is required"); + const Path path(source.filesystem().directory()); + if (!path.is_absolute() || std::filesystem::is_symlink(path) || !std::filesystem::is_directory(path)) + throw std::invalid_argument("source must be an absolute storage directory"); + return path; +} + +Path +DestinationPath(const StorageBackend& destination) +{ + if (!destination.has_filesystem() || destination.filesystem().directory().empty()) + throw std::invalid_argument("filesystem destination is required"); + const Path path(destination.filesystem().directory()); + if (!path.is_absolute()) + throw std::invalid_argument("destination must be an absolute storage directory"); + return path; +} + +Path +PartialPath(const Path& destination) +{ + Path partial = destination; + partial += ".pagebroker-partial"; + return partial; +} + +uintmax_t +DirectorySize(const Path& path) +{ + uintmax_t bytes = 0; + for (const auto& entry : std::filesystem::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; +} +} // namespace + TransferEngineType PosixCopyEngine::type() const { return TransferEngineType::POSIX_COPY; } +uintmax_t +PosixCopyEngine::RestoreSize(const StorageBackend& source) const +{ + return DirectorySize(SourcePath(source)); +} + +void +PosixCopyEngine::StageRestore(const StorageBackend& source, const Path& destination) const +{ + CopyDirectory(SourcePath(source), destination); +} + +bool +PosixCopyEngine::CheckpointDestinationConflicts(const StorageBackend& destination) const +{ + return std::filesystem::exists(PartialPath(DestinationPath(destination))); +} + +void +PosixCopyEngine::PublishCheckpoint(const Path& source, const StorageBackend& destination) const +{ + const Path published = DestinationPath(destination); + const Path partial = PartialPath(published); + try { + std::filesystem::create_directories(published.parent_path()); + CopyDirectory(source, partial); + std::filesystem::remove_all(published); + std::filesystem::rename(partial, published); + } + catch (...) { + std::error_code cleanup_error; + std::filesystem::remove_all(partial, cleanup_error); + throw; + } +} + void PosixCopyEngine::CopyDirectory(const Path& source, const Path& destination) const { diff --git a/deploy/snapshot/pagebroker/posix_copy_engine.hpp b/deploy/snapshot/pagebroker/posix_copy_engine.hpp index f46e88c766c3..601eabff491b 100644 --- a/deploy/snapshot/pagebroker/posix_copy_engine.hpp +++ b/deploy/snapshot/pagebroker/posix_copy_engine.hpp @@ -6,6 +6,10 @@ namespace snapshot::pagebroker { class PosixCopyEngine final : public TransferEngine { public: TransferEngineType type() const override; + uintmax_t RestoreSize(const StorageBackend& source) const override; + void StageRestore(const StorageBackend& source, const Path& destination) const override; + bool CheckpointDestinationConflicts(const StorageBackend& destination) const override; + void PublishCheckpoint(const Path& source, const StorageBackend& destination) const override; void CopyDirectory(const Path& source, const Path& destination) const override; }; } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/transaction.cpp b/deploy/snapshot/pagebroker/transaction.cpp index 9b0d5a937a2c..723abd6a290a 100644 --- a/deploy/snapshot/pagebroker/transaction.cpp +++ b/deploy/snapshot/pagebroker/transaction.cpp @@ -39,15 +39,4 @@ Transaction::clear_descriptor() descriptor_ = std::monostate(); } -uintmax_t -Transaction::reserved_bytes() const -{ - return reserved_bytes_; -} - -void -Transaction::set_reserved_bytes(uintmax_t bytes) -{ - reserved_bytes_ = bytes; -} } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/transaction.hpp b/deploy/snapshot/pagebroker/transaction.hpp index 7cf75bd0b617..b38a482bccc3 100644 --- a/deploy/snapshot/pagebroker/transaction.hpp +++ b/deploy/snapshot/pagebroker/transaction.hpp @@ -1,7 +1,6 @@ #pragma once #include -#include #include #include "checkpoint_transaction_descriptor.hpp" @@ -19,13 +18,10 @@ class Transaction { const Descriptor& descriptor() const; void set_descriptor(Descriptor descriptor); void clear_descriptor(); - uintmax_t reserved_bytes() const; - void set_reserved_bytes(uintmax_t bytes); private: std::mutex mutex_; State state_ = State::NEW; Descriptor descriptor_; - uintmax_t reserved_bytes_ = 0; }; } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/transfer_engine.hpp b/deploy/snapshot/pagebroker/transfer_engine.hpp index b4ea1165acce..e2731b56af6f 100644 --- a/deploy/snapshot/pagebroker/transfer_engine.hpp +++ b/deploy/snapshot/pagebroker/transfer_engine.hpp @@ -1,7 +1,10 @@ #pragma once +#include #include +#include "pagebroker_types.hpp" + namespace snapshot::pagebroker { using Path = std::filesystem::path; @@ -11,6 +14,10 @@ class TransferEngine { public: virtual ~TransferEngine(); virtual TransferEngineType type() const = 0; + virtual uintmax_t RestoreSize(const StorageBackend& source) const = 0; + virtual void StageRestore(const StorageBackend& source, const Path& destination) const = 0; + virtual bool CheckpointDestinationConflicts(const StorageBackend& destination) const = 0; + virtual void PublishCheckpoint(const Path& source, const StorageBackend& destination) const = 0; virtual void CopyDirectory(const Path& source, const Path& destination) const = 0; }; } // namespace snapshot::pagebroker From 6a2e2116130c9a058c9c6b5e24161c54865ea2d1 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 15:10:26 +0300 Subject: [PATCH 03/10] fix(snapshot): validate staged checkpoint destinations --- deploy/snapshot/pagebroker/broker.cpp | 1 + deploy/snapshot/pagebroker/daemon.cpp | 18 ++++++++++++++++-- .../snapshot/pagebroker/posix_copy_engine.cpp | 6 ++++++ .../snapshot/pagebroker/posix_copy_engine.hpp | 1 + deploy/snapshot/pagebroker/transfer_engine.hpp | 1 + 5 files changed, 25 insertions(+), 2 deletions(-) diff --git a/deploy/snapshot/pagebroker/broker.cpp b/deploy/snapshot/pagebroker/broker.cpp index 366fdf9f6414..2b790a8f33d5 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/pagebroker/broker.cpp @@ -247,6 +247,7 @@ Broker::StageCheckpoint(const Request& request, const StorageBackend& destinatio if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) return Fail(request, Failure::TRANSACTION_CONFLICT, "checkpoint transaction conflicts"); try { + engine.ValidateCheckpointDestination(destination); transaction.set_state(Transaction::State::PREPARING); fs::create_directory(staging_directory); transaction.set_descriptor(CheckpointTransactionDescriptor(staging_directory, destination, engine.type())); diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index bfeadb006603..6023ece62848 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -247,6 +247,20 @@ ServeConnection(int connection, Broker& broker) HandleConnection(descriptor.get(), broker); } +void +WaitForHandler(std::future& handler) +{ + try { + handler.get(); + } + catch (const std::exception& error) { + std::cerr << "connection handler: " << error.what() << '\n'; + } + catch (...) { + std::cerr << "connection handler: unknown exception\n"; + } +} + void ReapHandlers(std::vector>& handlers) { @@ -255,7 +269,7 @@ ReapHandlers(std::vector>& handlers) ++handler; continue; } - handler->get(); + WaitForHandler(*handler); handler = handlers.erase(handler); } } @@ -263,7 +277,7 @@ ReapHandlers(std::vector>& handlers) void WaitForHandlers(std::vector>& handlers) { - for (auto& handler : handlers) handler.get(); + for (auto& handler : handlers) WaitForHandler(handler); } void diff --git a/deploy/snapshot/pagebroker/posix_copy_engine.cpp b/deploy/snapshot/pagebroker/posix_copy_engine.cpp index 91dd2114d94e..bed21d7c7ab8 100644 --- a/deploy/snapshot/pagebroker/posix_copy_engine.cpp +++ b/deploy/snapshot/pagebroker/posix_copy_engine.cpp @@ -67,6 +67,12 @@ PosixCopyEngine::StageRestore(const StorageBackend& source, const Path& destinat CopyDirectory(SourcePath(source), destination); } +void +PosixCopyEngine::ValidateCheckpointDestination(const StorageBackend& destination) const +{ + DestinationPath(destination); +} + bool PosixCopyEngine::CheckpointDestinationConflicts(const StorageBackend& destination) const { diff --git a/deploy/snapshot/pagebroker/posix_copy_engine.hpp b/deploy/snapshot/pagebroker/posix_copy_engine.hpp index 601eabff491b..92fe68b5b22f 100644 --- a/deploy/snapshot/pagebroker/posix_copy_engine.hpp +++ b/deploy/snapshot/pagebroker/posix_copy_engine.hpp @@ -8,6 +8,7 @@ class PosixCopyEngine final : public TransferEngine { TransferEngineType type() const override; uintmax_t RestoreSize(const StorageBackend& source) const override; void StageRestore(const StorageBackend& source, const Path& destination) const override; + void ValidateCheckpointDestination(const StorageBackend& destination) const override; bool CheckpointDestinationConflicts(const StorageBackend& destination) const override; void PublishCheckpoint(const Path& source, const StorageBackend& destination) const override; void CopyDirectory(const Path& source, const Path& destination) const override; diff --git a/deploy/snapshot/pagebroker/transfer_engine.hpp b/deploy/snapshot/pagebroker/transfer_engine.hpp index e2731b56af6f..1cf2fc4f06bb 100644 --- a/deploy/snapshot/pagebroker/transfer_engine.hpp +++ b/deploy/snapshot/pagebroker/transfer_engine.hpp @@ -16,6 +16,7 @@ class TransferEngine { virtual TransferEngineType type() const = 0; virtual uintmax_t RestoreSize(const StorageBackend& source) const = 0; virtual void StageRestore(const StorageBackend& source, const Path& destination) const = 0; + virtual void ValidateCheckpointDestination(const StorageBackend& destination) const = 0; virtual bool CheckpointDestinationConflicts(const StorageBackend& destination) const = 0; virtual void PublishCheckpoint(const Path& source, const StorageBackend& destination) const = 0; virtual void CopyDirectory(const Path& source, const Path& destination) const = 0; From 4e5047e56f262a5d22e6493c0a218d4ad7ad9e78 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 15:27:07 +0300 Subject: [PATCH 04/10] fix(snapshot): validate PageBroker staging before state changes --- deploy/snapshot/pagebroker/broker.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deploy/snapshot/pagebroker/broker.cpp b/deploy/snapshot/pagebroker/broker.cpp index 2b790a8f33d5..1e470f94419f 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/pagebroker/broker.cpp @@ -202,11 +202,11 @@ Broker::StageRestore(const Request& request, const StorageBackend& source, const { const Path restore_root = staging_root_ / "restore"; const Path staging_directory = TransactionDirectory(restore_root, request.transaction_id()); - const uintmax_t bytes = engine.RestoreSize(source); Transaction& transaction = GetTransaction(request.transaction_id()); std::lock_guard lock(transaction.mutex()); if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) return Fail(request, Failure::TRANSACTION_CONFLICT, "restore transaction conflicts"); + const uintmax_t bytes = engine.RestoreSize(source); if (!ReserveStaging(bytes)) return Fail(request, Failure::INSUFFICIENT_STORAGE, "insufficient tmpfs capacity"); bool staging_reserved = true; @@ -242,12 +242,12 @@ Broker::StageCheckpoint(const Request& request, const StorageBackend& destinatio { const Path checkpoint_root = staging_root_ / "checkpoint"; const Path staging_directory = TransactionDirectory(checkpoint_root, request.transaction_id()); + engine.ValidateCheckpointDestination(destination); Transaction& transaction = GetTransaction(request.transaction_id()); std::lock_guard lock(transaction.mutex()); if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) return Fail(request, Failure::TRANSACTION_CONFLICT, "checkpoint transaction conflicts"); try { - engine.ValidateCheckpointDestination(destination); transaction.set_state(Transaction::State::PREPARING); fs::create_directory(staging_directory); transaction.set_descriptor(CheckpointTransactionDescriptor(staging_directory, destination, engine.type())); From 9f89e4ebd0a8d6ca24ff47549a37543f890d8cf2 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 15:46:15 +0300 Subject: [PATCH 05/10] fix(snapshot): bound PageBroker connection handlers --- .../charts/snapshot/templates/daemonset.yaml | 2 ++ deploy/helm/charts/snapshot/values.yaml | 2 ++ deploy/snapshot/pagebroker/daemon.cpp | 13 +++++++++---- deploy/snapshot/pagebroker/daemon.hpp | 4 +++- deploy/snapshot/pagebroker/main.cpp | 19 ++++++++++++++++--- 5 files changed, 32 insertions(+), 8 deletions(-) diff --git a/deploy/helm/charts/snapshot/templates/daemonset.yaml b/deploy/helm/charts/snapshot/templates/daemonset.yaml index d3dc5237177b..f1209ffd3d8d 100644 --- a/deploy/helm/charts/snapshot/templates/daemonset.yaml +++ b/deploy/helm/charts/snapshot/templates/daemonset.yaml @@ -113,6 +113,8 @@ spec: {{- else }} - {{ printf "%s/pagebroker.sock" $pageBrokerControlPath | quote }} - {{ $pageBrokerStagingPath | quote }} + - --max-concurrency + - {{ .Values.pageBroker.maxConcurrency | quote }} {{- end }} volumeMounts: - name: pagebroker diff --git a/deploy/helm/charts/snapshot/values.yaml b/deploy/helm/charts/snapshot/values.yaml index d9b0216dffc1..c6324064bcbf 100644 --- a/deploy/helm/charts/snapshot/values.yaml +++ b/deploy/helm/charts/snapshot/values.yaml @@ -84,6 +84,8 @@ pageBroker: command: [] # Optional arguments override. Defaults to the PageBroker socket and staging path. args: [] + # Maximum simultaneous PageBroker socket handlers. + maxConcurrency: 16 # Optional maximum size for tmpfs staging, for example "100Gi". stagingSizeLimit: "" image: diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index 6023ece62848..a0011154fe37 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -281,7 +281,7 @@ WaitForHandlers(std::vector>& handlers) } void -Serve(FileDescriptor& listener, Broker& broker) +Serve(FileDescriptor& listener, Broker& broker, size_t max_concurrency) { std::vector> handlers; while (!shutting_down) { @@ -301,11 +301,16 @@ Serve(FileDescriptor& listener, Broker& broker) LogError("accept", {errno, std::generic_category()}); continue; } + if (handlers.size() == max_concurrency) { + FileDescriptor descriptor(connection); + std::cerr << "connection limit reached\n"; + continue; + } try { handlers.emplace_back( std::async(std::launch::async, [connection, &broker] { ServeConnection(connection, broker); })); } - catch (const std::system_error& error) { + catch (const std::exception& error) { FileDescriptor descriptor(connection); std::cerr << "start connection: " << error.what() << '\n'; } @@ -315,7 +320,7 @@ Serve(FileDescriptor& listener, Broker& broker) } // namespace ExitCode -RunDaemon(const fs::path& socket_path, const fs::path& staging_directory) +RunDaemon(const fs::path& socket_path, const fs::path& staging_directory, size_t max_concurrency) { shutting_down = 0; if (const auto error = InstallSignalHandlers(); error) @@ -331,6 +336,6 @@ RunDaemon(const fs::path& socket_path, const fs::path& staging_directory) return Fail("create listener", error); Broker broker(staging_directory); - Serve(listener, broker); + Serve(listener, broker, max_concurrency); return ExitCode::SUCCESS; } diff --git a/deploy/snapshot/pagebroker/daemon.hpp b/deploy/snapshot/pagebroker/daemon.hpp index e9974ab9392b..3ef1a09f71d9 100644 --- a/deploy/snapshot/pagebroker/daemon.hpp +++ b/deploy/snapshot/pagebroker/daemon.hpp @@ -1,7 +1,9 @@ #pragma once +#include #include enum class ExitCode { SUCCESS = 0, FAILURE = 1, INVALID_ARGUMENTS = 2 }; -ExitCode RunDaemon(const std::filesystem::path& socket_path, const std::filesystem::path& staging_directory); +ExitCode RunDaemon( + const std::filesystem::path& socket_path, const std::filesystem::path& staging_directory, size_t max_concurrency); diff --git a/deploy/snapshot/pagebroker/main.cpp b/deploy/snapshot/pagebroker/main.cpp index 22f274f5967a..93fa646ce2c2 100644 --- a/deploy/snapshot/pagebroker/main.cpp +++ b/deploy/snapshot/pagebroker/main.cpp @@ -1,13 +1,26 @@ +#include #include +#include #include "daemon.hpp" +namespace { +bool +ParseMaxConcurrency(std::string_view value, size_t& max_concurrency) +{ + const auto [end, error] = std::from_chars(value.data(), value.data() + value.size(), max_concurrency); + return error == std::errc{} && end == value.data() + value.size() && max_concurrency > 0; +} +} // namespace + int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pagebroker-daemon SOCKET STAGING_DIRECTORY\n"; + size_t max_concurrency; + if (argc != 5 || std::string_view(argv[3]) != "--max-concurrency" || + !ParseMaxConcurrency(argv[4], max_concurrency)) { + std::cerr << "usage: pagebroker socket_path staging_directory --max-concurrency max_concurrency\n"; return static_cast(ExitCode::INVALID_ARGUMENTS); } - return static_cast(RunDaemon(argv[1], argv[2])); + return static_cast(RunDaemon(argv[1], argv[2], max_concurrency)); } From 632b1cce13920beeb007bae20e9d02649747e861 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 15:49:29 +0300 Subject: [PATCH 06/10] chore(compliance): attribute PageBroker in snapshot agent --- container/compliance/native_packages.yaml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/container/compliance/native_packages.yaml b/container/compliance/native_packages.yaml index 9b212475a57b..d25f503e2839 100644 --- a/container/compliance/native_packages.yaml +++ b/container/compliance/native_packages.yaml @@ -185,3 +185,12 @@ packages: source: https://github.com/ai-dynamo/dynamo images: - snapshot-agent + + - name: pagebroker + # First-party C++ daemon built from deploy/snapshot/pagebroker and copied + # into the snapshot-agent image. + version: 1.0 + license: Apache-2.0 + source: https://github.com/ai-dynamo/dynamo + images: + - snapshot-agent From 1b58a5d62e3ce912dab6fa1efb8282f716dd1d85 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 16:01:05 +0300 Subject: [PATCH 07/10] fix(snapshot): bound PageBroker transaction retention Signed-off-by: Dan Feigin --- .../charts/snapshot/templates/daemonset.yaml | 2 +- deploy/snapshot/internal/pagebroker/client.go | 11 +- deploy/snapshot/pagebroker/broker.cpp | 148 ++++++++++++------ deploy/snapshot/pagebroker/broker.hpp | 23 ++- deploy/snapshot/pagebroker/daemon_test.cpp | 71 +++++++++ deploy/snapshot/pagebroker/main.cpp | 3 +- deploy/snapshot/pagebroker/transaction.cpp | 9 ++ deploy/snapshot/pagebroker/transaction.hpp | 2 + 8 files changed, 216 insertions(+), 53 deletions(-) diff --git a/deploy/helm/charts/snapshot/templates/daemonset.yaml b/deploy/helm/charts/snapshot/templates/daemonset.yaml index f1209ffd3d8d..2181ca264e2f 100644 --- a/deploy/helm/charts/snapshot/templates/daemonset.yaml +++ b/deploy/helm/charts/snapshot/templates/daemonset.yaml @@ -114,7 +114,7 @@ spec: - {{ printf "%s/pagebroker.sock" $pageBrokerControlPath | quote }} - {{ $pageBrokerStagingPath | quote }} - --max-concurrency - - {{ .Values.pageBroker.maxConcurrency | quote }} + - {{ default 16 .Values.pageBroker.maxConcurrency | quote }} {{- end }} volumeMounts: - name: pagebroker diff --git a/deploy/snapshot/internal/pagebroker/client.go b/deploy/snapshot/internal/pagebroker/client.go index 7cad593dcde4..817ca63126bf 100644 --- a/deploy/snapshot/internal/pagebroker/client.go +++ b/deploy/snapshot/internal/pagebroker/client.go @@ -20,6 +20,7 @@ const ( // PageBroker control requests and responses are limited to 64 KiB. maxMessageSize = 64 << 10 commitRetryDelay = 100 * time.Millisecond + commitRetryLimit = 30 * time.Second ) var errMessageTooLarge = fmt.Errorf("message exceeds %d bytes", maxMessageSize) @@ -56,6 +57,7 @@ func (c Client) PrepareCheckpoint(ctx context.Context, transactionID, destinatio } func (c Client) Commit(ctx context.Context, transactionID string) error { + var retryDeadline time.Time for { response, err := c.request(ctx, transactionID, &Request_Commit{Commit: &CommitRequest{}}) if err == nil { @@ -68,10 +70,17 @@ func (c Client) Commit(ctx context.Context, transactionID string) error { if !errors.As(err, &transport) { return err } + if retryDeadline.IsZero() { + retryDeadline = time.Now().Add(commitRetryLimit) + } + delay := min(commitRetryDelay, time.Until(retryDeadline)) + if delay <= 0 { + return err + } select { case <-ctx.Done(): return ctx.Err() - case <-time.After(commitRetryDelay): + case <-time.After(delay): } } } diff --git a/deploy/snapshot/pagebroker/broker.cpp b/deploy/snapshot/pagebroker/broker.cpp index 1e470f94419f..022772906401 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/pagebroker/broker.cpp @@ -14,6 +14,9 @@ namespace snapshot::pagebroker { namespace fs = std::filesystem; namespace { +constexpr auto kTerminalTransactionRetention = std::chrono::seconds(60); +constexpr size_t kMaxRetainedTerminalTransactions = 1024; + Response Reply(const Request& request) { @@ -104,11 +107,56 @@ Broker::Broker(Path staging_root) : staging_root_(fs::weakly_canonical(std::move fs::create_directories(staging_root_ / "checkpoint"); } -Transaction& -Broker::GetTransaction(const std::string& transaction_id) +Broker::TransactionHandle +Broker::CreateOrGetTransaction(const std::string& transaction_id) +{ + std::lock_guard lock(transactions_mutex_); + auto [iterator, inserted] = transactions_.try_emplace(transaction_id, std::make_shared()); + return iterator->second; +} + +Broker::TransactionHandle +Broker::FindTransaction(const std::string& transaction_id) +{ + std::lock_guard lock(transactions_mutex_); + const auto iterator = transactions_.find(transaction_id); + return iterator == transactions_.end() ? nullptr : iterator->second; +} + +void +Broker::RetainTerminalTransaction(const std::string& transaction_id) +{ + auto transaction = FindTransaction(transaction_id); + if (!transaction) + return; + std::lock_guard transaction_lock(transaction->mutex()); + if (!transaction->retain_terminal()) + return; + std::lock_guard terminal_lock(terminal_transactions_mutex_); + terminal_transactions_.push_back({transaction_id, std::move(transaction), std::chrono::steady_clock::now()}); +} + +void +Broker::ReapTerminalTransactions() { + const auto now = std::chrono::steady_clock::now(); + std::vector expired; + { + std::lock_guard lock(terminal_transactions_mutex_); + while (!terminal_transactions_.empty() && + (now - terminal_transactions_.front().completed >= kTerminalTransactionRetention || + terminal_transactions_.size() > kMaxRetainedTerminalTransactions)) { + expired.push_back(std::move(terminal_transactions_.front())); + terminal_transactions_.pop_front(); + } + } + std::lock_guard lock(transactions_mutex_); - return transactions_.try_emplace(transaction_id).first->second; + for (const auto& item : expired) { + const auto iterator = transactions_.find(item.id); + if (iterator != transactions_.end() && iterator->second == item.transaction) + transactions_.erase(iterator); + } } bool @@ -167,18 +215,26 @@ Broker::HandleRequest(const Request& request) return Fail(request, Failure::INVALID_REQUEST, "request and transaction IDs are required"); try { + Response response; switch (request.command_case()) { case Request::kStagedRestore: - return Restore(request); + response = Restore(request); + break; case Request::kPrepareStagedCheckpoint: - return PrepareCheckpoint(request); + response = PrepareCheckpoint(request); + break; case Request::kCommit: - return Commit(request); + response = Commit(request); + break; case Request::kAbort: - return Abort(request); + response = Abort(request); + break; default: return Fail(request, Failure::INVALID_REQUEST, "unsupported operation"); } + RetainTerminalTransaction(request.transaction_id()); + ReapTerminalTransactions(); + return response; } catch (const std::invalid_argument& error) { return Fail(request, Failure::INVALID_REQUEST, error.what()); @@ -202,26 +258,26 @@ Broker::StageRestore(const Request& request, const StorageBackend& source, const { const Path restore_root = staging_root_ / "restore"; const Path staging_directory = TransactionDirectory(restore_root, request.transaction_id()); - Transaction& transaction = GetTransaction(request.transaction_id()); - std::lock_guard lock(transaction.mutex()); - if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) - return Fail(request, Failure::TRANSACTION_CONFLICT, "restore transaction conflicts"); const uintmax_t bytes = engine.RestoreSize(source); + auto transaction = CreateOrGetTransaction(request.transaction_id()); + std::lock_guard lock(transaction->mutex()); + if (transaction->state() != Transaction::State::NEW || fs::exists(staging_directory)) + return Fail(request, Failure::TRANSACTION_CONFLICT, "restore transaction conflicts"); if (!ReserveStaging(bytes)) return Fail(request, Failure::INSUFFICIENT_STORAGE, "insufficient tmpfs capacity"); bool staging_reserved = true; try { - transaction.set_state(Transaction::State::PREPARING); + transaction->set_state(Transaction::State::PREPARING); engine.StageRestore(source, staging_directory); ReleaseStaging(bytes); staging_reserved = false; - transaction.set_descriptor(RestoreTransactionDescriptor(staging_directory)); - transaction.set_state(Transaction::State::STAGED); + transaction->set_descriptor(RestoreTransactionDescriptor(staging_directory)); + transaction->set_state(Transaction::State::STAGED); } catch (const std::exception& error) { if (staging_reserved) ReleaseStaging(bytes); - return AbortStaging(request, transaction, staging_directory, error); + return AbortStaging(request, *transaction, staging_directory, error); } auto response = Reply(request); response.mutable_staged_restore_directory()->set_image_directory(staging_directory.string()); @@ -243,18 +299,18 @@ Broker::StageCheckpoint(const Request& request, const StorageBackend& destinatio const Path checkpoint_root = staging_root_ / "checkpoint"; const Path staging_directory = TransactionDirectory(checkpoint_root, request.transaction_id()); engine.ValidateCheckpointDestination(destination); - Transaction& transaction = GetTransaction(request.transaction_id()); - std::lock_guard lock(transaction.mutex()); - if (transaction.state() != Transaction::State::NEW || fs::exists(staging_directory)) + auto transaction = CreateOrGetTransaction(request.transaction_id()); + std::lock_guard lock(transaction->mutex()); + if (transaction->state() != Transaction::State::NEW || fs::exists(staging_directory)) return Fail(request, Failure::TRANSACTION_CONFLICT, "checkpoint transaction conflicts"); try { - transaction.set_state(Transaction::State::PREPARING); + transaction->set_state(Transaction::State::PREPARING); fs::create_directory(staging_directory); - transaction.set_descriptor(CheckpointTransactionDescriptor(staging_directory, destination, engine.type())); - transaction.set_state(Transaction::State::STAGED); + transaction->set_descriptor(CheckpointTransactionDescriptor(staging_directory, destination, engine.type())); + transaction->set_state(Transaction::State::STAGED); } catch (const std::exception& error) { - return AbortStaging(request, transaction, staging_directory, error); + return AbortStaging(request, *transaction, staging_directory, error); } auto response = Reply(request); response.mutable_staged_checkpoint_directory()->set_image_directory(staging_directory.string()); @@ -264,37 +320,38 @@ Broker::StageCheckpoint(const Request& request, const StorageBackend& destinatio Response Broker::Commit(const Request& request) { - Transaction& transaction = GetTransaction(request.transaction_id()); - std::lock_guard lock(transaction.mutex()); - if (transaction.state() == Transaction::State::NEW || transaction.state() == Transaction::State::ABORTED) + auto transaction = FindTransaction(request.transaction_id()); + if (!transaction) return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); - if (transaction.state() == Transaction::State::PREPARING) + std::lock_guard lock(transaction->mutex()); + if (transaction->state() == Transaction::State::NEW || transaction->state() == Transaction::State::ABORTED) + return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); + if (transaction->state() == Transaction::State::PREPARING) return Fail(request, Failure::TRANSACTION_CONFLICT, "transaction is preparing"); - if (transaction.state() == Transaction::State::COMMITTED) + if (transaction->state() == Transaction::State::COMMITTED) return CommitSucceeded(request); - if (const auto* restore = std::get_if(&transaction.descriptor())) - return CleanupRestore(request, *restore); + if (const auto* restore = std::get_if(&transaction->descriptor())) + return CleanupRestore(request, *transaction, *restore); - const auto* checkpoint = std::get_if(&transaction.descriptor()); + const auto* checkpoint = std::get_if(&transaction->descriptor()); if (checkpoint == nullptr) return Fail(request, Failure::INTERNAL_ERROR, "live transaction has no descriptor"); - - return PublishCheckpoint(request, *checkpoint); + return PublishCheckpoint(request, *transaction, *checkpoint); } Response -Broker::CleanupRestore(const Request& request, const RestoreTransactionDescriptor& transaction) +Broker::CleanupRestore(const Request& request, Transaction& transaction, const RestoreTransactionDescriptor& descriptor) { - fs::remove_all(transaction.staging_directory()); - Transaction& current = GetTransaction(request.transaction_id()); - current.clear_descriptor(); - current.set_state(Transaction::State::COMMITTED); + fs::remove_all(descriptor.staging_directory()); + transaction.clear_descriptor(); + transaction.set_state(Transaction::State::COMMITTED); return CommitSucceeded(request); } Response -Broker::PublishCheckpoint(const Request& request, const CheckpointTransactionDescriptor& transaction) +Broker::PublishCheckpoint( + const Request& request, Transaction& current, const CheckpointTransactionDescriptor& transaction) { const Path staging_directory = transaction.staging_directory(); if (!fs::is_directory(staging_directory)) @@ -305,7 +362,6 @@ Broker::PublishCheckpoint(const Request& request, const CheckpointTransactionDes RejectSymlinks(staging_directory); try { engine.PublishCheckpoint(staging_directory, transaction.destination_storage()); - Transaction& current = GetTransaction(request.transaction_id()); current.clear_descriptor(); current.set_state(Transaction::State::COMMITTED); std::error_code cleanup_error; @@ -320,19 +376,21 @@ Broker::PublishCheckpoint(const Request& request, const CheckpointTransactionDes Response Broker::Abort(const Request& request) { - Transaction& transaction = GetTransaction(request.transaction_id()); - std::lock_guard lock(transaction.mutex()); - if (transaction.state() == Transaction::State::NEW || transaction.state() == Transaction::State::COMMITTED) + auto transaction = FindTransaction(request.transaction_id()); + if (!transaction) return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); - if (transaction.state() == Transaction::State::ABORTED) + std::lock_guard lock(transaction->mutex()); + if (transaction->state() == Transaction::State::NEW || transaction->state() == Transaction::State::COMMITTED) + return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); + if (transaction->state() == Transaction::State::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())); - transaction.clear_descriptor(); - transaction.set_state(Transaction::State::ABORTED); + transaction->clear_descriptor(); + transaction->set_state(Transaction::State::ABORTED); return AbortSucceeded(request); } diff --git a/deploy/snapshot/pagebroker/broker.hpp b/deploy/snapshot/pagebroker/broker.hpp index 16c64fd85390..fddfe4c018f2 100644 --- a/deploy/snapshot/pagebroker/broker.hpp +++ b/deploy/snapshot/pagebroker/broker.hpp @@ -1,5 +1,7 @@ #pragma once +#include +#include #include #include #include @@ -20,11 +22,20 @@ class Broker { private: using Engines = std::vector>; - using Transactions = std::unordered_map; + using TransactionHandle = std::shared_ptr; + using Transactions = std::unordered_map; + struct TerminalTransaction { + std::string id; + TransactionHandle transaction; + std::chrono::steady_clock::time_point completed; + }; const TransferEngine& Engine(TransferEngineType engine_type) const; const TransferEngine& Engine(const IOEngine& engine) const; - Transaction& GetTransaction(const std::string& transaction_id); + TransactionHandle CreateOrGetTransaction(const std::string& transaction_id); + TransactionHandle FindTransaction(const std::string& transaction_id); + void RetainTerminalTransaction(const std::string& transaction_id); + void ReapTerminalTransactions(); bool ReserveStaging(uintmax_t bytes); void ReleaseStaging(uintmax_t bytes); Response AbortStaging( @@ -35,13 +46,17 @@ class Broker { Response StageCheckpoint(const Request& request, const StorageBackend& destination, const TransferEngine& engine); // 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 CleanupRestore( + const Request& request, Transaction& transaction, const RestoreTransactionDescriptor& descriptor); + Response PublishCheckpoint( + const Request& request, Transaction& transaction, const CheckpointTransactionDescriptor& descriptor); Response Abort(const Request& request); Path staging_root_; Engines io_engines_; std::mutex transactions_mutex_; Transactions transactions_; + std::mutex terminal_transactions_mutex_; + std::deque terminal_transactions_; uintmax_t reserved_staging_bytes_ = 0; }; } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/daemon_test.cpp b/deploy/snapshot/pagebroker/daemon_test.cpp index 338209beb139..cd69ad955cd5 100644 --- a/deploy/snapshot/pagebroker/daemon_test.cpp +++ b/deploy/snapshot/pagebroker/daemon_test.cpp @@ -145,6 +145,21 @@ TEST_F(BrokerTest, RejectsSymlinkInRestoreSource) EXPECT_EQ(response.failure().code(), Failure::STORAGE_ERROR); } +TEST_F(BrokerTest, InvalidRestoreDoesNotReserveTransaction) +{ + auto invalid = RequestFor("restore"); + Configure( + invalid.mutable_staged_restore()->mutable_source(), invalid.mutable_staged_restore()->mutable_io_engine(), + "relative"); + EXPECT_EQ(broker().HandleRequest(invalid).failure().code(), Failure::INVALID_REQUEST); + + auto restore = RequestFor("restore"); + Configure( + restore.mutable_staged_restore()->mutable_source(), restore.mutable_staged_restore()->mutable_io_engine(), + source_); + EXPECT_TRUE(broker().HandleRequest(restore).has_staged_restore_directory()); +} + TEST_F(BrokerTest, RejectsInvalidStagedRestore) { auto restore = RequestFor("restore"); @@ -154,6 +169,62 @@ TEST_F(BrokerTest, RejectsInvalidStagedRestore) EXPECT_EQ(response.failure().code(), Failure::INVALID_REQUEST); } +TEST_F(BrokerTest, UnknownCommitAndAbortDoNotReserveTransactions) +{ + auto commit = RequestFor("unknown-commit"); + commit.mutable_commit(); + EXPECT_EQ(broker().HandleRequest(commit).failure().code(), Failure::TRANSACTION_NOT_FOUND); + + auto restore = RequestFor("unknown-commit"); + Configure( + restore.mutable_staged_restore()->mutable_source(), restore.mutable_staged_restore()->mutable_io_engine(), + source_); + EXPECT_TRUE(broker().HandleRequest(restore).has_staged_restore_directory()); + + auto abort = RequestFor("unknown-abort"); + abort.mutable_abort(); + EXPECT_EQ(broker().HandleRequest(abort).failure().code(), Failure::TRANSACTION_NOT_FOUND); + + auto prepare = RequestFor("unknown-abort"); + Configure( + prepare.mutable_prepare_staged_checkpoint()->mutable_destination(), + prepare.mutable_prepare_staged_checkpoint()->mutable_io_engine(), root_ / "storage" / "published"); + EXPECT_TRUE(broker().HandleRequest(prepare).has_staged_checkpoint_directory()); +} + +TEST_F(BrokerTest, EvictsOldTerminalTransactionsButRetainsRecentCompletions) +{ + auto oldest = RequestFor("oldest"); + Configure( + oldest.mutable_prepare_staged_checkpoint()->mutable_destination(), + oldest.mutable_prepare_staged_checkpoint()->mutable_io_engine(), root_ / "storage" / "oldest"); + ASSERT_TRUE(broker().HandleRequest(oldest).has_staged_checkpoint_directory()); + auto oldest_commit = RequestFor("oldest"); + oldest_commit.mutable_commit(); + ASSERT_TRUE(broker().HandleRequest(oldest_commit).has_commit_complete()); + + for (size_t index = 0; index < 1024; ++index) { + const auto id = "terminal-" + std::to_string(index); + auto prepare = RequestFor(id); + Configure( + prepare.mutable_prepare_staged_checkpoint()->mutable_destination(), + prepare.mutable_prepare_staged_checkpoint()->mutable_io_engine(), root_ / "storage" / id); + ASSERT_TRUE(broker().HandleRequest(prepare).has_staged_checkpoint_directory()); + auto commit = RequestFor(id); + commit.mutable_commit(); + ASSERT_TRUE(broker().HandleRequest(commit).has_commit_complete()); + } + + auto recent_commit = RequestFor("terminal-1023"); + recent_commit.mutable_commit(); + EXPECT_TRUE(broker().HandleRequest(recent_commit).has_commit_complete()); + + auto reuse = RequestFor("oldest"); + Configure( + reuse.mutable_staged_restore()->mutable_source(), reuse.mutable_staged_restore()->mutable_io_engine(), source_); + EXPECT_TRUE(broker().HandleRequest(reuse).has_staged_restore_directory()); +} + TEST_F(BrokerTest, AbortsFailedCheckpointStaging) { const fs::path destination = root_ / "storage" / "published"; diff --git a/deploy/snapshot/pagebroker/main.cpp b/deploy/snapshot/pagebroker/main.cpp index 93fa646ce2c2..6dae503d94ea 100644 --- a/deploy/snapshot/pagebroker/main.cpp +++ b/deploy/snapshot/pagebroker/main.cpp @@ -17,8 +17,7 @@ int main(int argc, char** argv) { size_t max_concurrency; - if (argc != 5 || std::string_view(argv[3]) != "--max-concurrency" || - !ParseMaxConcurrency(argv[4], max_concurrency)) { + if (argc != 5 || std::string_view(argv[3]) != "--max-concurrency" || !ParseMaxConcurrency(argv[4], max_concurrency)) { std::cerr << "usage: pagebroker socket_path staging_directory --max-concurrency max_concurrency\n"; return static_cast(ExitCode::INVALID_ARGUMENTS); } diff --git a/deploy/snapshot/pagebroker/transaction.cpp b/deploy/snapshot/pagebroker/transaction.cpp index 723abd6a290a..9f24d3007b29 100644 --- a/deploy/snapshot/pagebroker/transaction.cpp +++ b/deploy/snapshot/pagebroker/transaction.cpp @@ -39,4 +39,13 @@ Transaction::clear_descriptor() descriptor_ = std::monostate(); } +bool +Transaction::retain_terminal() +{ + if (terminal_retained_ || (state_ != State::COMMITTED && state_ != State::ABORTED)) + return false; + terminal_retained_ = true; + return true; +} + } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/transaction.hpp b/deploy/snapshot/pagebroker/transaction.hpp index b38a482bccc3..e042ff761f58 100644 --- a/deploy/snapshot/pagebroker/transaction.hpp +++ b/deploy/snapshot/pagebroker/transaction.hpp @@ -18,10 +18,12 @@ class Transaction { const Descriptor& descriptor() const; void set_descriptor(Descriptor descriptor); void clear_descriptor(); + bool retain_terminal(); private: std::mutex mutex_; State state_ = State::NEW; Descriptor descriptor_; + bool terminal_retained_ = false; }; } // namespace snapshot::pagebroker From f99de7cebe6a1af640d9bc62482657dba8fea7f4 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 18:30:38 +0300 Subject: [PATCH 08/10] feat(snapshot): time PageBroker restore phases --- deploy/snapshot/internal/executor/restore.go | 45 ++++++++++++++------ deploy/snapshot/pagebroker/daemon.cpp | 5 ++- 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/deploy/snapshot/internal/executor/restore.go b/deploy/snapshot/internal/executor/restore.go index 691cfc360642..afe7fe629f2b 100644 --- a/deploy/snapshot/internal/executor/restore.go +++ b/deploy/snapshot/internal/executor/restore.go @@ -66,6 +66,10 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r transactionID := "" var broker pagebroker.Client committed := false + var pageBrokerStageDuration time.Duration + var pageBrokerStagingMountDuration time.Duration + var pageBrokerStagingUnmountDuration time.Duration + var pageBrokerCommitDuration time.Duration brokered := req.PageBrokerRequested && req.PageBrokerEnabled if brokered { transactionID = uuid.NewString() @@ -77,7 +81,9 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r _ = broker.Abort(abortCtx, transactionID) } }() + stageStart := time.Now() staged, err := broker.StagedRestore(ctx, transactionID, req.CheckpointLocation) + pageBrokerStageDuration = time.Since(stageStart) if err != nil { return 0, fmt.Errorf("stage PageBroker restore: %w", err) } @@ -100,6 +106,7 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r } injectDuration := time.Since(injectStart) defer func() { + unmountStart := time.Now() // Pass a background context: mp.Unmount has its own internal timeout // (nsmount.unmountTimeout) around the ns-bind-mount subprocess. if cleanupErr := mp.Unmount(context.Background()); cleanupErr != nil { @@ -109,10 +116,12 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r // already restored successfully. Log it and let the pod continue. log.Error(cleanupErr, "failed to unmount agent bundle from placeholder namespace") } + log.Info("Agent bundle unmount timing", "duration", time.Since(unmountStart)) }() var mountedStaging nsmount.MountPoint if brokered { + stagingMountStart := time.Now() stagingMounter, err := nsmount.New(req.CheckpointLocation, nsmount.PageBrokerDst, log) if err != nil { return 0, fmt.Errorf("create PageBroker staging mount: %w", err) @@ -121,36 +130,51 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r if err != nil { return 0, fmt.Errorf("mount PageBroker staging: %w", err) } + pageBrokerStagingMountDuration = time.Since(stagingMountStart) req.ContainerCheckpointLocation = nsmount.PageBrokerDst } // Phase 3: Execute — nsrestore handles rootfs, CRIU restore, and CUDA restore inside namespace. result, err := execNSRestore(ctx, log, req, snap, mp) if mountedStaging != nil { + stagingUnmountStart := time.Now() if cleanupErr := mountedStaging.Unmount(context.Background()); cleanupErr != nil { log.Error(cleanupErr, "failed to unmount PageBroker staging from placeholder namespace") } + pageBrokerStagingUnmountDuration = time.Since(stagingUnmountStart) } if err != nil { return 0, fmt.Errorf("nsrestore failed: %w", err) } if brokered { + commitStart := time.Now() if err := broker.Commit(ctx, transactionID); err != nil { log.Error(err, "failed to commit PageBroker restore") } else { committed = true } + pageBrokerCommitDuration = time.Since(commitStart) } - restoreDuration := hostInspectDuration + injectDuration + result.TotalDuration() + + validationStart := time.Now() + if err := validateRestoredProcess(snap.TargetRoot, result.RestoredPID, log); err != nil { + return 0, err + } + validationDuration := time.Since(validationStart) log.Info("Restore timing summary", "restore", map[string]any{ - "duration": restoreDuration.String(), + "duration": time.Since(restoreStart).String(), "phases": map[string]string{ - "host_inspect_duration": hostInspectDuration.String(), - "inject_duration": injectDuration.String(), - "nsrestore_setup_duration": result.NSRestoreSetupDuration.String(), - "criu_restore_duration": result.CRIURestoreDuration.String(), - "cuda_duration": result.CUDADuration.String(), + "pagebroker_stage_duration": pageBrokerStageDuration.String(), + "host_inspect_duration": hostInspectDuration.String(), + "inject_duration": injectDuration.String(), + "pagebroker_staging_mount_duration": pageBrokerStagingMountDuration.String(), + "nsrestore_setup_duration": result.NSRestoreSetupDuration.String(), + "criu_restore_duration": result.CRIURestoreDuration.String(), + "cuda_duration": result.CUDADuration.String(), + "pagebroker_staging_unmount_duration": pageBrokerStagingUnmountDuration.String(), + "pagebroker_commit_duration": pageBrokerCommitDuration.String(), + "validation_duration": validationDuration.String(), }, }, ) @@ -160,15 +184,10 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r ) } - validationStart := time.Now() - if err := validateRestoredProcess(snap.TargetRoot, result.RestoredPID, log); err != nil { - return 0, err - } - log.Info("=== External restore completed ===", "restored_pid", result.RestoredPID, "placeholder_host_pid", snap.PlaceholderPID, - "validation_duration", time.Since(validationStart), + "validation_duration", validationDuration, "total_duration", time.Since(restoreStart), ) diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index a0011154fe37..fea42373bc50 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -222,10 +222,13 @@ HandleConnection(int connection, Broker& broker) if (!ReadAll(connection, message.data(), size) || !request.ParseFromString(message) || !request.IsInitialized()) { response = InvalidRequest(); } else { + const auto request_start = std::chrono::steady_clock::now(); response = broker.HandleRequest(request); + const auto duration = + std::chrono::duration_cast(std::chrono::steady_clock::now() - request_start); std::osyncstream(std::cerr) << "transaction=" << request.transaction_id() << " command=" << CommandName(request.command_case()) - << " result=" << ResultName(response) + << " result=" << ResultName(response) << " duration_ms=" << duration.count() << (response.has_failure() ? " error=" + response.failure().message() : "") << '\n'; } } From 348e41df86938b36412abf39ed5535c1c38c3cff Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Sun, 16 Aug 2026 14:30:32 +0300 Subject: [PATCH 09/10] fix(snapshot): expire PageBroker staging --- deploy/snapshot/pagebroker/broker.cpp | 33 ++++++++++++++++++++++ deploy/snapshot/pagebroker/broker.hpp | 1 + deploy/snapshot/pagebroker/daemon.cpp | 7 +++++ deploy/snapshot/pagebroker/daemon_test.cpp | 32 +++++++++++++++++++++ deploy/snapshot/pagebroker/transaction.cpp | 8 ++++++ deploy/snapshot/pagebroker/transaction.hpp | 3 ++ 6 files changed, 84 insertions(+) diff --git a/deploy/snapshot/pagebroker/broker.cpp b/deploy/snapshot/pagebroker/broker.cpp index 022772906401..9872097f297b 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/pagebroker/broker.cpp @@ -16,6 +16,7 @@ namespace { constexpr auto kTerminalTransactionRetention = std::chrono::seconds(60); constexpr size_t kMaxRetainedTerminalTransactions = 1024; +constexpr auto kLiveTransactionLifetime = std::chrono::hours(1); Response Reply(const Request& request) @@ -103,10 +104,42 @@ TransactionDirectory(const Path& transaction_root, const std::string& transactio Broker::Broker(Path staging_root) : staging_root_(fs::weakly_canonical(std::move(staging_root))) { io_engines_.push_back(std::make_unique()); + fs::remove_all(staging_root_ / "restore"); + fs::remove_all(staging_root_ / "checkpoint"); fs::create_directories(staging_root_ / "restore"); fs::create_directories(staging_root_ / "checkpoint"); } +void +Broker::ReapExpiredTransactions(std::chrono::steady_clock::time_point now) +{ + std::vector> transactions; + { + std::lock_guard lock(transactions_mutex_); + for (const auto& [id, transaction] : transactions_) transactions.emplace_back(id, transaction); + } + + for (const auto& [id, transaction] : transactions) { + std::lock_guard transaction_lock(transaction->mutex()); + if (!transaction->expired(now, kLiveTransactionLifetime)) + continue; + + std::error_code restore_error; + std::error_code checkpoint_error; + fs::remove_all(TransactionDirectory(staging_root_ / "restore", id), restore_error); + fs::remove_all(TransactionDirectory(staging_root_ / "checkpoint", id), checkpoint_error); + if (restore_error || checkpoint_error) + continue; + transaction->clear_descriptor(); + transaction->set_state(Transaction::State::ABORTED); + + std::lock_guard transactions_lock(transactions_mutex_); + const auto current = transactions_.find(id); + if (current != transactions_.end() && current->second == transaction) + transactions_.erase(current); + } +} + Broker::TransactionHandle Broker::CreateOrGetTransaction(const std::string& transaction_id) { diff --git a/deploy/snapshot/pagebroker/broker.hpp b/deploy/snapshot/pagebroker/broker.hpp index fddfe4c018f2..f61c6b685cd9 100644 --- a/deploy/snapshot/pagebroker/broker.hpp +++ b/deploy/snapshot/pagebroker/broker.hpp @@ -19,6 +19,7 @@ class Broker { public: explicit Broker(Path staging_root); Response HandleRequest(const Request& request); + void ReapExpiredTransactions(std::chrono::steady_clock::time_point now); private: using Engines = std::vector>; diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index fea42373bc50..36bfde9cb2d4 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -33,6 +33,7 @@ namespace { constexpr uint32_t kMaxMessageSize = 64 << 10; // 64 KB constexpr timeval kConnectionTimeout{30, 0}; constexpr int kShutdownPollTimeoutMs = 1000; +constexpr auto kTransactionReapInterval = std::chrono::minutes(2); volatile sig_atomic_t shutting_down; void @@ -287,8 +288,14 @@ void Serve(FileDescriptor& listener, Broker& broker, size_t max_concurrency) { std::vector> handlers; + auto next_transaction_reap = std::chrono::steady_clock::now(); while (!shutting_down) { ReapHandlers(handlers); + const auto now = std::chrono::steady_clock::now(); + if (now >= next_transaction_reap) { + broker.ReapExpiredTransactions(now); + next_transaction_reap = now + kTransactionReapInterval; + } pollfd poll_descriptor{listener.get(), POLLIN, 0}; const int ready = poll(&poll_descriptor, 1, kShutdownPollTimeoutMs); if (ready == 0) diff --git a/deploy/snapshot/pagebroker/daemon_test.cpp b/deploy/snapshot/pagebroker/daemon_test.cpp index cd69ad955cd5..c9d05e642939 100644 --- a/deploy/snapshot/pagebroker/daemon_test.cpp +++ b/deploy/snapshot/pagebroker/daemon_test.cpp @@ -120,6 +120,38 @@ TEST_F(BrokerTest, RejectsConcurrentRestoreForSameTransaction) EXPECT_NE(first_response.has_staged_restore_directory(), second_response.has_staged_restore_directory()); } +TEST_F(BrokerTest, ReapsExpiredStagedTransactions) +{ + auto restore = RequestFor("expired"); + Configure( + restore.mutable_staged_restore()->mutable_source(), restore.mutable_staged_restore()->mutable_io_engine(), source_); + const auto staged = broker().HandleRequest(restore); + ASSERT_TRUE(staged.has_staged_restore_directory()); + const fs::path staging_directory(staged.staged_restore_directory().image_directory()); + + broker().ReapExpiredTransactions(std::chrono::steady_clock::now() + std::chrono::hours(1)); + EXPECT_FALSE(fs::exists(staging_directory)); + + auto commit = RequestFor("expired"); + commit.mutable_commit(); + EXPECT_EQ(broker().HandleRequest(commit).failure().code(), Failure::TRANSACTION_NOT_FOUND); + + auto retry = RequestFor("expired"); + Configure(retry.mutable_staged_restore()->mutable_source(), retry.mutable_staged_restore()->mutable_io_engine(), source_); + EXPECT_TRUE(broker().HandleRequest(retry).has_staged_restore_directory()); +} + +TEST_F(BrokerTest, CleansStaleStagingOnStart) +{ + broker_.reset(); + const fs::path stale = root_ / "tmpfs" / "restore" / "stale"; + fs::create_directories(stale); + std::ofstream(stale / "image") << "image"; + + broker_.emplace(root_ / "tmpfs"); + EXPECT_FALSE(fs::exists(stale)); +} + TEST_F(BrokerTest, RejectsUnsafeTransactionIDs) { for (const auto& id : diff --git a/deploy/snapshot/pagebroker/transaction.cpp b/deploy/snapshot/pagebroker/transaction.cpp index 9f24d3007b29..b251ede0f603 100644 --- a/deploy/snapshot/pagebroker/transaction.cpp +++ b/deploy/snapshot/pagebroker/transaction.cpp @@ -19,6 +19,8 @@ void Transaction::set_state(State state) { state_ = state; + if (state == State::PREPARING) + staging_started_at_ = std::chrono::steady_clock::now(); } const Transaction::Descriptor& @@ -48,4 +50,10 @@ Transaction::retain_terminal() return true; } +bool +Transaction::expired(std::chrono::steady_clock::time_point now, std::chrono::steady_clock::duration lifetime) const +{ + return state_ == State::STAGED && now - staging_started_at_ >= lifetime; +} + } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/transaction.hpp b/deploy/snapshot/pagebroker/transaction.hpp index e042ff761f58..10308e84dc18 100644 --- a/deploy/snapshot/pagebroker/transaction.hpp +++ b/deploy/snapshot/pagebroker/transaction.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -19,11 +20,13 @@ class Transaction { void set_descriptor(Descriptor descriptor); void clear_descriptor(); bool retain_terminal(); + bool expired(std::chrono::steady_clock::time_point now, std::chrono::steady_clock::duration lifetime) const; private: std::mutex mutex_; State state_ = State::NEW; Descriptor descriptor_; + std::chrono::steady_clock::time_point staging_started_at_; bool terminal_retained_ = false; }; } // namespace snapshot::pagebroker From 9d1d91094e9f24fc887537f87e9471a138521fd6 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Sun, 16 Aug 2026 16:28:51 +0300 Subject: [PATCH 10/10] fix(snapshot): bound PageBroker terminal replies --- deploy/snapshot/pagebroker/broker.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deploy/snapshot/pagebroker/broker.cpp b/deploy/snapshot/pagebroker/broker.cpp index 9872097f297b..ea003eca79e3 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/pagebroker/broker.cpp @@ -14,7 +14,7 @@ namespace snapshot::pagebroker { namespace fs = std::filesystem; namespace { -constexpr auto kTerminalTransactionRetention = std::chrono::seconds(60); +constexpr auto kTerminalTransactionRetention = std::chrono::hours(1); constexpr size_t kMaxRetainedTerminalTransactions = 1024; constexpr auto kLiveTransactionLifetime = std::chrono::hours(1);