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 diff --git a/deploy/helm/charts/snapshot/templates/daemonset.yaml b/deploy/helm/charts/snapshot/templates/daemonset.yaml index d3dc5237177b..2181ca264e2f 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 + - {{ default 16 .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/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/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/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..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 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) 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 8c10c27388ee..ea003eca79e3 100644 --- a/deploy/snapshot/pagebroker/broker.cpp +++ b/deploy/snapshot/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 @@ -94,10 +104,124 @@ 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) +{ + 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) +{ + 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 { @@ -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) { @@ -116,18 +248,29 @@ 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()); } catch (const std::exception& error) { return Fail(request, Failure::STORAGE_ERROR, error.what()); @@ -138,30 +281,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"); - 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)) 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 +321,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 +353,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& current, const CheckpointTransactionDescriptor& transaction) { const Path staging_directory = transaction.staging_directory(); if (!fs::is_directory(staging_directory)) return Fail(request, Failure::TRANSACTION_NOT_FOUND, "checkpoint staging directory not found"); - const Path published_directory(transaction.destination_storage().filesystem().directory()); - Path partial = published_directory; - partial += ".pagebroker-partial"; - if (fs::exists(partial)) + 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); - checkpoint_transactions_.erase(request.transaction_id()); - transaction_states_.at(request.transaction_id()) = TransactionState::COMMITTED; + engine.PublishCheckpoint(staging_directory, transaction.destination_storage()); + 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); } @@ -263,19 +409,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/deploy/snapshot/pagebroker/broker.hpp b/deploy/snapshot/pagebroker/broker.hpp index 2fde5f25229f..f61c6b685cd9 100644 --- a/deploy/snapshot/pagebroker/broker.hpp +++ b/deploy/snapshot/pagebroker/broker.hpp @@ -1,12 +1,17 @@ #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 { @@ -14,26 +19,45 @@ 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>; - 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/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index 4131aed609ea..36bfde9cb2d4 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/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,113 @@ 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_concurrency) +{ + 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_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::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, size_t max_concurrency) +{ + 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, max_concurrency); + return ExitCode::SUCCESS; } diff --git a/deploy/snapshot/pagebroker/daemon.hpp b/deploy/snapshot/pagebroker/daemon.hpp new file mode 100644 index 000000000000..3ef1a09f71d9 --- /dev/null +++ b/deploy/snapshot/pagebroker/daemon.hpp @@ -0,0 +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, size_t max_concurrency); diff --git a/deploy/snapshot/pagebroker/daemon_test.cpp b/deploy/snapshot/pagebroker/daemon_test.cpp index a6da7ca2c7af..c9d05e642939 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,80 @@ 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(); + + 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 : @@ -102,6 +177,108 @@ 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"); + 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/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..6dae503d94ea --- /dev/null +++ b/deploy/snapshot/pagebroker/main.cpp @@ -0,0 +1,25 @@ +#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) +{ + 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], max_concurrency)); +} diff --git a/deploy/snapshot/pagebroker/posix_copy_engine.cpp b/deploy/snapshot/pagebroker/posix_copy_engine.cpp index 44be57509ba1..bed21d7c7ab8 100644 --- a/deploy/snapshot/pagebroker/posix_copy_engine.cpp +++ b/deploy/snapshot/pagebroker/posix_copy_engine.cpp @@ -1,14 +1,102 @@ #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); +} + +void +PosixCopyEngine::ValidateCheckpointDestination(const StorageBackend& destination) const +{ + DestinationPath(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..92fe68b5b22f 100644 --- a/deploy/snapshot/pagebroker/posix_copy_engine.hpp +++ b/deploy/snapshot/pagebroker/posix_copy_engine.hpp @@ -6,6 +6,11 @@ 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; + 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; }; } // namespace snapshot::pagebroker diff --git a/deploy/snapshot/pagebroker/transaction.cpp b/deploy/snapshot/pagebroker/transaction.cpp new file mode 100644 index 000000000000..b251ede0f603 --- /dev/null +++ b/deploy/snapshot/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/deploy/snapshot/pagebroker/transaction.hpp b/deploy/snapshot/pagebroker/transaction.hpp new file mode 100644 index 000000000000..10308e84dc18 --- /dev/null +++ b/deploy/snapshot/pagebroker/transaction.hpp @@ -0,0 +1,32 @@ +#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(); + 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/deploy/snapshot/pagebroker/transfer_engine.hpp b/deploy/snapshot/pagebroker/transfer_engine.hpp index b4ea1165acce..1cf2fc4f06bb 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,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