From f1bc573d369272afbc63fba8fb66df6f14461595 Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 20 Aug 2026 12:38:53 +0300 Subject: [PATCH] fix(agent): bound PageBroker terminal replies --- agent/internal/executor/restore.go | 13 + agent/internal/pagebroker/client.go | 11 +- agent/internal/types/config.go | 3 + agent/internal/types/config_test.go | 9 + agent/pagebroker/Makefile | 4 +- agent/pagebroker/broker.cpp | 330 ++++++++++++++++++------- agent/pagebroker/broker.hpp | 44 +++- agent/pagebroker/daemon.cpp | 284 +++++++++++++++++---- agent/pagebroker/daemon.hpp | 12 + agent/pagebroker/daemon_test.cpp | 236 +++++++++++++++++- agent/pagebroker/file_descriptor.cpp | 15 ++ agent/pagebroker/file_descriptor.hpp | 2 + agent/pagebroker/main.cpp | 26 ++ agent/pagebroker/posix_copy_engine.cpp | 103 ++++++++ agent/pagebroker/posix_copy_engine.hpp | 9 + agent/pagebroker/transaction.cpp | 59 +++++ agent/pagebroker/transaction.hpp | 33 +++ agent/pagebroker/transfer_engine.hpp | 8 + 18 files changed, 1053 insertions(+), 148 deletions(-) create mode 100644 agent/pagebroker/daemon.hpp create mode 100644 agent/pagebroker/main.cpp create mode 100644 agent/pagebroker/transaction.cpp create mode 100644 agent/pagebroker/transaction.hpp diff --git a/agent/internal/executor/restore.go b/agent/internal/executor/restore.go index 483a953e..6de23e6e 100644 --- a/agent/internal/executor/restore.go +++ b/agent/internal/executor/restore.go @@ -151,6 +151,7 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r transactionID := "" var broker pagebroker.Client committed := false + var pageBrokerStageDuration, pageBrokerMountDuration, pageBrokerCommitDuration time.Duration if brokered { transactionID = uuid.NewString() broker = pagebroker.Client{ControlSocketPath: req.PageBrokerControlSocketPath} @@ -161,11 +162,15 @@ 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, artifactPath) + pageBrokerStageDuration = time.Since(stageStart) if err != nil { return 0, fmt.Errorf("stage PageBroker restore: %w", err) } + mountStart := time.Now() stagingMount, err := mounts.MountPageBroker(ctx, bundleMount, staged) + pageBrokerMountDuration = time.Since(mountStart) if err != nil { return 0, fmt.Errorf("mount PageBroker staging: %w", err) } @@ -190,11 +195,13 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r 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) } if result.CleanupError != nil { cleanupErr = errors.Join(cleanupErr, result.CleanupError) @@ -206,6 +213,9 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r cleanup() wall := time.Since(restoreStart) unaccounted := remainingDuration(wall, + pageBrokerStageDuration, + pageBrokerMountDuration, + pageBrokerCommitDuration, gpuDeviceMapDuration, result.OverlayCaptureDuration, result.CRIUPrepareDuration, @@ -215,6 +225,9 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r summary := map[string]any{ "duration": wall.String(), "phases": map[string]string{ + "pagebroker_stage": pageBrokerStageDuration.String(), + "pagebroker_mount": pageBrokerMountDuration.String(), + "pagebroker_commit": pageBrokerCommitDuration.String(), "gpu_device_map": gpuDeviceMapDuration.String(), "overlay_capture": result.OverlayCaptureDuration.String(), "criu_prepare": result.CRIUPrepareDuration.String(), diff --git a/agent/internal/pagebroker/client.go b/agent/internal/pagebroker/client.go index 1fb7d86e..fd75582d 100644 --- a/agent/internal/pagebroker/client.go +++ b/agent/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) @@ -57,6 +58,7 @@ func imageDirectory(directory string) (string, error) { } 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 { @@ -69,10 +71,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/agent/internal/types/config.go b/agent/internal/types/config.go index 86bbce38..e0185887 100644 --- a/agent/internal/types/config.go +++ b/agent/internal/types/config.go @@ -45,6 +45,9 @@ func (c *AgentConfig) Validate() error { return &ConfigError{Field: "storage.basePath", Message: fmt.Sprintf("storage.basePath must be %q", CheckpointBasePath)} } c.Storage.BasePath = basePath + 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/agent/internal/types/config_test.go b/agent/internal/types/config_test.go index 27e1af81..ff7cfa39 100644 --- a/agent/internal/types/config_test.go +++ b/agent/internal/types/config_test.go @@ -26,3 +26,12 @@ func TestAgentConfigValidateRequiresFixedStorageBasePath(t *testing.T) { } } } + +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/agent/pagebroker/Makefile b/agent/pagebroker/Makefile index 4ea687d2..3cfea0ef 100644 --- a/agent/pagebroker/Makefile +++ b/agent/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 +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) daemon.cpp main.cpp file_descriptor.cpp .PHONY: daemon generate test diff --git a/agent/pagebroker/broker.cpp b/agent/pagebroker/broker.cpp index 8c10c273..c8c64bd3 100644 --- a/agent/pagebroker/broker.cpp +++ b/agent/pagebroker/broker.cpp @@ -14,6 +14,10 @@ namespace snapshot::pagebroker { namespace fs = std::filesystem; namespace { +constexpr auto kTerminalTransactionRetention = std::chrono::hours(1); +constexpr size_t kMaxRetainedTerminalTransactions = 1024; +constexpr auto kLiveTransactionLifetime = std::chrono::hours(1); + Response Reply(const Request& request) { @@ -55,23 +59,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) +{ + if (!request.has_source() || request.source().kind_case() == StorageBackend::KIND_NOT_SET) + throw std::invalid_argument("restore source is required"); + return request.source(); +} + +const StorageBackend& +ValidateStagedCheckpoint(const PrepareStagedCheckpointRequest& request) { - return storage.has_filesystem() && !storage.filesystem().directory().empty() && engine.has_posix_copy(); + if (!request.has_destination() || request.destination().kind_case() == StorageBackend::KIND_NOT_SET) + throw std::invalid_argument("checkpoint destination is required"); + return request.destination(); } -uintmax_t -TreeSize(const Path& path) +void +RejectSymlinks(const Path& directory) { - uintmax_t bytes = 0; - for (const auto& entry : fs::recursive_directory_iterator(path)) { + 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 @@ -91,13 +101,127 @@ TransactionDirectory(const Path& transaction_root, const std::string& transactio } // namespace -Broker::Broker(Path staging_root) : staging_root_(fs::weakly_canonical(std::move(staging_root))) +Broker::Broker(Path staging_root, Path storage_root) : staging_root_(fs::weakly_canonical(std::move(staging_root))) { - io_engines_.push_back(std::make_unique()); + io_engines_.push_back(std::make_unique(std::move(storage_root))); + 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) +{ + 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_); + for (const auto& item : expired) { + const auto iterator = transactions_.find(item.id); + if (iterator != transactions_.end() && iterator->second == item.transaction) + transactions_.erase(iterator); + } +} + +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; +} + +Response +Broker::AbortStaging( + const Request& request, Transaction& transaction, const Path& staging_directory, const std::exception& error) +{ + transaction.clear_descriptor(); + transaction.set_state(Transaction::State::ABORTED); + 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()); + return Fail(request, Failure::STORAGE_ERROR, error.what()); +} + const TransferEngine& Broker::Engine(TransferEngineType engine_type) const { @@ -108,6 +232,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) { @@ -115,53 +247,76 @@ Broker::HandleRequest(const Request& request) !IsSafePathComponent(request.transaction_id())) return Fail(request, Failure::INVALID_REQUEST, "request and transaction IDs are required"); + Response response; try { 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"); + response = Fail(request, Failure::INVALID_REQUEST, "unsupported operation"); + break; } } + catch (const std::invalid_argument& error) { + response = Fail(request, Failure::INVALID_REQUEST, error.what()); + } catch (const std::exception& error) { - return Fail(request, Failure::STORAGE_ERROR, error.what()); + response = Fail(request, Failure::STORAGE_ERROR, error.what()); } + RetainTerminalTransaction(request.transaction_id()); + ReapTerminalTransactions(); + return response; } 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"); - if (fs::exists(staging_directory) || transaction_states_.contains(request.transaction_id())) + 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 (!HasAvailableSpace(staging_root_, TreeSize(source))) + if (!ReserveStaging(bytes)) { + std::lock_guard transactions_lock(transactions_mutex_); + const auto current = transactions_.find(request.transaction_id()); + if (current != transactions_.end() && current->second == transaction) + transactions_.erase(current); return Fail(request, Failure::INSUFFICIENT_STORAGE, "insufficient tmpfs capacity"); - + } + bool staging_reserved = true; try { - transaction_states_.emplace(request.transaction_id(), TransactionState::LIVE); - engine.CopyDirectory(source, staging_directory); - restore_transactions_.emplace(request.transaction_id(), RestoreTransactionDescriptor(staging_directory)); + 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); } - catch (...) { - fs::remove_all(staging_directory); - transaction_states_.erase(request.transaction_id()); - restore_transactions_.erase(request.transaction_id()); - 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()); @@ -172,29 +327,29 @@ 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"); - if (fs::exists(staging_directory) || transaction_states_.contains(request.transaction_id())) + engine.ValidateCheckpointDestination(destination); + 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_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, 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()); - 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()); @@ -204,58 +359,55 @@ 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) + auto transaction = FindTransaction(request.transaction_id()); + if (!transaction) + return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); + 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, *transaction, *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, *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()); - restore_transactions_.erase(request.transaction_id()); - transaction_states_.at(request.transaction_id()) = TransactionState::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& transaction, const CheckpointTransactionDescriptor& descriptor) { - const Path staging_directory = transaction.staging_directory(); + const Path staging_directory = descriptor.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(descriptor.engine_type()); + if (engine.CheckpointDestinationConflicts(descriptor.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); - checkpoint_transactions_.erase(request.transaction_id()); - transaction_states_.at(request.transaction_id()) = TransactionState::COMMITTED; + engine.PublishCheckpoint(staging_directory, descriptor.destination_storage()); + transaction.clear_descriptor(); + transaction.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); } @@ -263,19 +415,21 @@ 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) + auto transaction = FindTransaction(request.transaction_id()); + if (!transaction) + return Fail(request, Failure::TRANSACTION_NOT_FOUND, "transaction not found"); + 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/agent/pagebroker/broker.hpp b/agent/pagebroker/broker.hpp index 2fde5f25..1efdc85f 100644 --- a/agent/pagebroker/broker.hpp +++ b/agent/pagebroker/broker.hpp @@ -1,39 +1,63 @@ #pragma once +#include +#include +#include #include +#include #include #include #include "checkpoint_transaction_descriptor.hpp" #include "pagebroker_types.hpp" #include "restore_transaction_descriptor.hpp" +#include "transaction.hpp" #include "transfer_engine.hpp" namespace snapshot::pagebroker { class Broker { public: - explicit Broker(Path staging_root); + Broker(Path staging_root, Path storage_root); Response HandleRequest(const Request& request); + void ReapExpiredTransactions(std::chrono::steady_clock::time_point now); 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 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; + 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( + 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); - 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_; - CheckpointTransactions checkpoint_transactions_; - RestoreTransactions restore_transactions_; - TransactionStates transaction_states_; + 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/agent/pagebroker/daemon.cpp b/agent/pagebroker/daemon.cpp index 4131aed6..a25e281c 100644 --- a/agent/pagebroker/daemon.cpp +++ b/agent/pagebroker/daemon.cpp @@ -1,13 +1,24 @@ +#include "daemon.hpp" + #include #include +#include +#include #include #include #include +#include #include #include +#include #include #include +#include +#include +#include +#include +#include #include "broker.hpp" #include "file_descriptor.hpp" @@ -20,7 +31,9 @@ using snapshot::pagebroker::Response; namespace { constexpr uint32_t kMaxMessageSize = 64 << 10; // 64 KB -enum ArgumentIndex { kSocketPath = 1, kStagingDirectory, kArgumentCount }; +constexpr timeval kConnectionTimeout{30, 0}; +constexpr int kShutdownPollTimeoutMs = 1000; +constexpr auto kTransactionReapInterval = std::chrono::minutes(2); volatile sig_atomic_t shutting_down; void @@ -29,6 +42,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) { @@ -74,6 +170,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) { @@ -91,7 +223,14 @@ 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) << " duration_ms=" << duration.count() + << (response.has_failure() ? " error=" + response.failure().message() : "") << '\n'; } } @@ -100,64 +239,117 @@ HandleConnection(int connection, Broker& broker) WriteAll(connection, &size, sizeof(size)); WriteAll(connection, message.data(), message.size()); } -} // namespace -int -main(int argc, char** argv) +void +ServeConnection(int connection, Broker& broker) { - 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; + FileDescriptor descriptor(connection); + if (const auto error = ConfigureConnection(descriptor.get()); error) { + LogError("set connection timeout", error); + return; } + HandleConnection(descriptor.get(), broker); +} - 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; +void +WaitForHandler(std::future& handler) +{ + try { + handler.get(); } - fs::create_directories(argv[kStagingDirectory], error); - if (error) { - std::cerr << "create staging directory: " << error.message() << '\n'; - return 1; + catch (const std::exception& error) { + std::cerr << "connection handler: " << error.what() << '\n'; } - if (socket_path.string().size() >= sizeof(sockaddr_un::sun_path)) { - std::cerr << "socket path is too long\n"; - return 2; + catch (...) { + std::cerr << "connection handler: unknown exception\n"; } - 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; +void +ReapHandlers(std::vector>& handlers) +{ + for (auto handler = handlers.begin(); handler != handlers.end();) { + if (handler->wait_for(std::chrono::seconds(0)) != std::future_status::ready) { + ++handler; + continue; + } + WaitForHandler(*handler); + handler = handlers.erase(handler); } +} + +void +WaitForHandlers(std::vector>& handlers) +{ + for (auto& handler : handlers) WaitForHandler(handler); +} - Broker broker(argv[kStagingDirectory]); +void +Serve(FileDescriptor& listener, Broker& broker, size_t max_concurrent_requests) +{ + std::vector> handlers; + auto next_transaction_reap = std::chrono::steady_clock::now(); while (!shutting_down) { - FileDescriptor connection(accept(listener.get(), nullptr, nullptr)); - if (connection.get() < 0) { + 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) + continue; + if (ready < 0) { if (errno != EINTR) - std::cerr << "accept: " << std::strerror(errno) << '\n'; + LogError("poll", {errno, std::generic_category()}); continue; } - HandleConnection(connection.get(), broker); + const int connection = accept(listener.get(), nullptr, nullptr); + if (connection < 0) { + if (errno != EINTR && errno != EAGAIN && errno != EWOULDBLOCK) + LogError("accept", {errno, std::generic_category()}); + continue; + } + if (handlers.size() == max_concurrent_requests) { + 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::exception& error) { + FileDescriptor descriptor(connection); + std::cerr << "start connection: " << error.what() << '\n'; + } } + WaitForHandlers(handlers); +} +} // namespace + +ExitCode +RunDaemon( + const fs::path& socket_path, + const fs::path& staging_directory, + const fs::path& storage_root, + size_t max_concurrent_requests) +{ + 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, storage_root); + Serve(listener, broker, max_concurrent_requests); + return ExitCode::SUCCESS; } diff --git a/agent/pagebroker/daemon.hpp b/agent/pagebroker/daemon.hpp new file mode 100644 index 00000000..26d1f2e5 --- /dev/null +++ b/agent/pagebroker/daemon.hpp @@ -0,0 +1,12 @@ +#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, + const std::filesystem::path& storage_root, + size_t max_concurrent_requests); diff --git a/agent/pagebroker/daemon_test.cpp b/agent/pagebroker/daemon_test.cpp index a6da7ca2..3b6b6449 100644 --- a/agent/pagebroker/daemon_test.cpp +++ b/agent/pagebroker/daemon_test.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "broker.hpp" @@ -22,7 +23,7 @@ class BrokerTest : public ::testing::Test { source_ = root_ / "storage" / "source"; fs::create_directories(source_); std::ofstream(source_ / "image") << "image"; - broker_.emplace(root_ / "tmpfs"); + broker_.emplace(root_ / "tmpfs", root_ / "storage"); } void TearDown() override { fs::remove_all(root_); } @@ -77,6 +78,83 @@ 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, 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(); + + ASSERT_NE(first_response.has_staged_restore_directory(), second_response.has_staged_restore_directory()); + const auto& rejected = + first_response.has_staged_restore_directory() ? second_response : first_response; + EXPECT_EQ(rejected.failure().code(), Failure::TRANSACTION_CONFLICT); +} + +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", root_ / "storage"); + EXPECT_FALSE(fs::exists(stale)); +} + TEST_F(BrokerTest, RejectsUnsafeTransactionIDs) { for (const auto& id : @@ -102,6 +180,162 @@ 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, RejectsPathsOutsideStorageRoot) +{ + const fs::path outside = root_ / "outside"; + fs::create_directories(outside); + std::ofstream(outside / "image") << "image"; + + auto restore = RequestFor("outside-restore"); + Configure( + restore.mutable_staged_restore()->mutable_source(), restore.mutable_staged_restore()->mutable_io_engine(), outside); + EXPECT_EQ(broker().HandleRequest(restore).failure().code(), Failure::INVALID_REQUEST); + + auto checkpoint = RequestFor("outside-checkpoint"); + Configure( + checkpoint.mutable_prepare_staged_checkpoint()->mutable_destination(), + checkpoint.mutable_prepare_staged_checkpoint()->mutable_io_engine(), + outside / "checkpoint"); + EXPECT_EQ(broker().HandleRequest(checkpoint).failure().code(), Failure::INVALID_REQUEST); + + fs::create_directory_symlink(outside, root_ / "storage" / "link"); + auto symlinked = RequestFor("symlinked-checkpoint"); + Configure( + symlinked.mutable_prepare_staged_checkpoint()->mutable_destination(), + symlinked.mutable_prepare_staged_checkpoint()->mutable_io_engine(), + root_ / "storage" / "link" / "checkpoint"); + EXPECT_EQ(broker().HandleRequest(symlinked).failure().code(), Failure::INVALID_REQUEST); + + auto root_destination = RequestFor("root-destination"); + Configure( + root_destination.mutable_prepare_staged_checkpoint()->mutable_destination(), + root_destination.mutable_prepare_staged_checkpoint()->mutable_io_engine(), root_ / "storage"); + EXPECT_EQ(broker().HandleRequest(root_destination).failure().code(), Failure::INVALID_REQUEST); +} + +TEST_F(BrokerTest, InsufficientStagingDoesNotReserveTransaction) +{ + const fs::path large = root_ / "storage" / "large"; + fs::create_directories(large); + std::ofstream file(large / "image"); + file.seekp(1LL << 40); + file.put('\0'); + file.close(); + + auto insufficient = RequestFor("restore"); + Configure( + insufficient.mutable_staged_restore()->mutable_source(), insufficient.mutable_staged_restore()->mutable_io_engine(), + large); + EXPECT_EQ(broker().HandleRequest(insufficient).failure().code(), Failure::INSUFFICIENT_STORAGE); + + auto retry = RequestFor("restore"); + 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, 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, 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"; + 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/agent/pagebroker/file_descriptor.cpp b/agent/pagebroker/file_descriptor.cpp index 88993926..074adae2 100644 --- a/agent/pagebroker/file_descriptor.cpp +++ b/agent/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/agent/pagebroker/file_descriptor.hpp b/agent/pagebroker/file_descriptor.hpp index b178cf54..c858259a 100644 --- a/agent/pagebroker/file_descriptor.hpp +++ b/agent/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/agent/pagebroker/main.cpp b/agent/pagebroker/main.cpp new file mode 100644 index 00000000..56d044dd --- /dev/null +++ b/agent/pagebroker/main.cpp @@ -0,0 +1,26 @@ +#include +#include +#include + +#include "daemon.hpp" + +namespace { +bool +ParseMaxConcurrentRequests(std::string_view value, size_t& max_concurrent_requests) +{ + const auto [end, error] = std::from_chars(value.data(), value.data() + value.size(), max_concurrent_requests); + return error == std::errc{} && end == value.data() + value.size() && max_concurrent_requests > 0; +} +} // namespace + +int +main(int argc, char** argv) +{ + size_t max_concurrent_requests; + if (argc != 6 || std::string_view(argv[4]) != "--max-concurrent-requests" || + !ParseMaxConcurrentRequests(argv[5], max_concurrent_requests)) { + std::cerr << "usage: pagebroker socket_path staging_directory storage_root --max-concurrent-requests max_concurrent_requests\n"; + return static_cast(ExitCode::INVALID_ARGUMENTS); + } + return static_cast(RunDaemon(argv[1], argv[2], argv[3], max_concurrent_requests)); +} diff --git a/agent/pagebroker/posix_copy_engine.cpp b/agent/pagebroker/posix_copy_engine.cpp index 44be5750..8485fd2d 100644 --- a/agent/pagebroker/posix_copy_engine.cpp +++ b/agent/pagebroker/posix_copy_engine.cpp @@ -1,14 +1,117 @@ #include "posix_copy_engine.hpp" #include +#include namespace snapshot::pagebroker { +namespace { +Path +StoragePath(const StorageBackend& storage, const Path& storage_root, const char* label) +{ + if (!storage.has_filesystem() || storage.filesystem().directory().empty()) + throw std::invalid_argument(std::string("filesystem ") + label + " is required"); + const Path path(storage.filesystem().directory()); + const Path relative = path.lexically_relative(storage_root); + if (!path.is_absolute() || path.lexically_normal() != path || relative.empty() || + relative == "." || relative.string().starts_with("../") || relative == "..") + throw std::invalid_argument(std::string(label) + " must be within storage root"); + + Path component = storage_root; + for (const auto& part : relative) { + component /= part; + if (std::filesystem::is_symlink(component)) + throw std::invalid_argument(std::string(label) + " contains symlink"); + } + return path; +} + +Path +SourcePath(const StorageBackend& source, const Path& storage_root) +{ + const Path path = StoragePath(source, storage_root, "source"); + if (!std::filesystem::is_directory(path)) + throw std::invalid_argument("source must be a storage directory"); + return path; +} + +Path +DestinationPath(const StorageBackend& destination, const Path& storage_root) +{ + return StoragePath(destination, storage_root, "destination"); +} + +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 + +PosixCopyEngine::PosixCopyEngine(Path storage_root) : storage_root_(std::filesystem::weakly_canonical(std::move(storage_root))) {} + TransferEngineType PosixCopyEngine::type() const { return TransferEngineType::POSIX_COPY; } +uintmax_t +PosixCopyEngine::RestoreSize(const StorageBackend& source) const +{ + return DirectorySize(SourcePath(source, storage_root_)); +} + +void +PosixCopyEngine::StageRestore(const StorageBackend& source, const Path& destination) const +{ + CopyDirectory(SourcePath(source, storage_root_), destination); +} + +void +PosixCopyEngine::ValidateCheckpointDestination(const StorageBackend& destination) const +{ + DestinationPath(destination, storage_root_); +} + +bool +PosixCopyEngine::CheckpointDestinationConflicts(const StorageBackend& destination) const +{ + return std::filesystem::exists(PartialPath(DestinationPath(destination, storage_root_))); +} + +void +PosixCopyEngine::PublishCheckpoint(const Path& source, const StorageBackend& destination) const +{ + const Path published = DestinationPath(destination, storage_root_); + 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/agent/pagebroker/posix_copy_engine.hpp b/agent/pagebroker/posix_copy_engine.hpp index f46e88c7..7cb61712 100644 --- a/agent/pagebroker/posix_copy_engine.hpp +++ b/agent/pagebroker/posix_copy_engine.hpp @@ -5,7 +5,16 @@ namespace snapshot::pagebroker { class PosixCopyEngine final : public TransferEngine { public: + explicit PosixCopyEngine(Path storage_root); 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; + + private: + Path storage_root_; }; } // namespace snapshot::pagebroker diff --git a/agent/pagebroker/transaction.cpp b/agent/pagebroker/transaction.cpp new file mode 100644 index 00000000..b251ede0 --- /dev/null +++ b/agent/pagebroker/transaction.cpp @@ -0,0 +1,59 @@ +#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; + if (state == State::PREPARING) + staging_started_at_ = std::chrono::steady_clock::now(); +} + +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(); +} + +bool +Transaction::retain_terminal() +{ + if (terminal_retained_ || (state_ != State::COMMITTED && state_ != State::ABORTED)) + return false; + terminal_retained_ = true; + 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/agent/pagebroker/transaction.hpp b/agent/pagebroker/transaction.hpp new file mode 100644 index 00000000..290a008c --- /dev/null +++ b/agent/pagebroker/transaction.hpp @@ -0,0 +1,33 @@ +#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; + + // Callers hold mutex() while accessing transaction state. + std::mutex& mutex(); + State state() const; + void set_state(State state); + const Descriptor& descriptor() const; + 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 diff --git a/agent/pagebroker/transfer_engine.hpp b/agent/pagebroker/transfer_engine.hpp index b4ea1165..1cf2fc4f 100644 --- a/agent/pagebroker/transfer_engine.hpp +++ b/agent/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,11 @@ 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 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; }; } // namespace snapshot::pagebroker