From 2aaad9ef73315bc13429b67b790374c789103f5d Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 10:51:23 +0300 Subject: [PATCH 1/2] feat(snapshot): add PageBroker Unix socket server --- deploy/snapshot/pagebroker/Dockerfile | 19 +++ deploy/snapshot/pagebroker/Makefile | 11 +- deploy/snapshot/pagebroker/daemon.cpp | 146 ++++++++++++++++++ .../snapshot/pagebroker/file_descriptor.cpp | 17 ++ .../snapshot/pagebroker/file_descriptor.hpp | 15 ++ 5 files changed, 205 insertions(+), 3 deletions(-) create mode 100644 deploy/snapshot/pagebroker/Dockerfile create mode 100644 deploy/snapshot/pagebroker/daemon.cpp create mode 100644 deploy/snapshot/pagebroker/file_descriptor.cpp create mode 100644 deploy/snapshot/pagebroker/file_descriptor.hpp diff --git a/deploy/snapshot/pagebroker/Dockerfile b/deploy/snapshot/pagebroker/Dockerfile new file mode 100644 index 000000000000..ceddb2047c66 --- /dev/null +++ b/deploy/snapshot/pagebroker/Dockerfile @@ -0,0 +1,19 @@ +FROM ubuntu:24.04 AS build + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + libprotobuf-dev \ + protobuf-compiler \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /src +COPY . . +RUN make daemon + +FROM ubuntu:24.04 + +RUN apt-get update && apt-get install -y --no-install-recommends libprotobuf32 \ + && rm -rf /var/lib/apt/lists/* + +COPY --from=build /src/pagebroker /usr/local/bin/pagebroker +ENTRYPOINT ["/usr/local/bin/pagebroker"] diff --git a/deploy/snapshot/pagebroker/Makefile b/deploy/snapshot/pagebroker/Makefile index ba4b42f55e28..4ea687d22933 100644 --- a/deploy/snapshot/pagebroker/Makefile +++ b/deploy/snapshot/pagebroker/Makefile @@ -1,12 +1,17 @@ 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 -.PHONY: generate test +.PHONY: daemon generate test generate: protoc --proto_path=. --cpp_out=. $(PROTO) -test: generate broker.cpp daemon_test.cpp +daemon: generate + c++ -I. -std=c++20 -Wall -Werror $(DAEMON_SOURCES) v1/pagebroker.pb.cc -lprotobuf -o pagebroker + +test: generate $(BROKER_SOURCES) daemon_test.cpp mkdir -p build - c++ -I. -std=c++20 -Wall -Werror broker.cpp daemon_test.cpp v1/pagebroker.pb.cc -lprotobuf $(GTEST_FLAGS) -o build/pagebroker-test + c++ -I. -std=c++20 -Wall -Werror $(BROKER_SOURCES) daemon_test.cpp v1/pagebroker.pb.cc -lprotobuf $(GTEST_FLAGS) -o build/pagebroker-test ./build/pagebroker-test diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp new file mode 100644 index 000000000000..fd8357397349 --- /dev/null +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -0,0 +1,146 @@ +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "broker.hpp" +#include "file_descriptor.hpp" + +namespace fs = std::filesystem; +using snapshot::pagebroker::Broker; +using snapshot::pagebroker::Failure; +using snapshot::pagebroker::Request; +using snapshot::pagebroker::Response; + +namespace { +constexpr uint32_t kMaxMessageSize = 64 << 10; // 64 KB +enum ArgumentIndex { kSocketPath = 1, kStagingDirectory, kArgumentCount }; + +bool +ReadAll(int fd, void* buffer, size_t size) +{ + auto* bytes = static_cast(buffer); + while (size > 0) { + ssize_t read; + do { + read = recv(fd, bytes, size, 0); + } while (read < 0 && errno == EINTR); + if (read <= 0) + return false; + bytes += read; + size -= read; + } + return true; +} + +bool +WriteAll(int fd, const void* buffer, size_t size) +{ + const auto* bytes = static_cast(buffer); + while (size > 0) { + ssize_t written; + do { + written = send(fd, bytes, size, MSG_NOSIGNAL); + } while (written < 0 && errno == EINTR); + if (written <= 0) + return false; + bytes += written; + size -= written; + } + return true; +} + +Response +InvalidRequest() +{ + Response response; + response.set_request_id(""); + response.set_transaction_id(""); + response.mutable_failure()->set_code(Failure::INVALID_REQUEST); + response.mutable_failure()->set_message("invalid request"); + return response; +} + +void +HandleConnection(int connection, Broker& broker) +{ + uint32_t size = 0; + if (!ReadAll(connection, &size, sizeof(size))) + return; + size = ntohl(size); + + Response response; + if (size > kMaxMessageSize) { + response = InvalidRequest(); + } else { + std::string message(size, '\0'); + Request request; + if (!ReadAll(connection, message.data(), size) || !request.ParseFromString(message) || !request.IsInitialized()) { + response = InvalidRequest(); + } else { + response = broker.HandleRequest(request); + } + } + + std::string message = response.SerializeAsString(); + size = htonl(message.size()); + WriteAll(connection, &size, sizeof(size)); + WriteAll(connection, message.data(), message.size()); +} +} // namespace + +int +main(int argc, char** argv) +{ + if (argc != kArgumentCount) { + std::cerr << "usage: pagebroker-daemon SOCKET STAGING_DIRECTORY\n"; + return 2; + } + + const fs::path socket_path(argv[kSocketPath]); + std::error_code error; + fs::create_directories(socket_path.parent_path(), error); + if (error) { + std::cerr << "create socket directory: " << error.message() << '\n'; + return 1; + } + fs::create_directories(argv[kStagingDirectory], error); + if (error) { + std::cerr << "create staging directory: " << error.message() << '\n'; + return 1; + } + if (socket_path.string().size() >= sizeof(sockaddr_un::sun_path)) { + std::cerr << "socket path is too long\n"; + return 2; + } + 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; + } + + Broker broker(argv[kStagingDirectory]); + for (;;) { + FileDescriptor connection(accept(listener.get(), nullptr, nullptr)); + if (connection.get() < 0) { + if (errno != EINTR) + std::cerr << "accept: " << std::strerror(errno) << '\n'; + continue; + } + HandleConnection(connection.get(), broker); + } +} diff --git a/deploy/snapshot/pagebroker/file_descriptor.cpp b/deploy/snapshot/pagebroker/file_descriptor.cpp new file mode 100644 index 000000000000..88993926c7f3 --- /dev/null +++ b/deploy/snapshot/pagebroker/file_descriptor.cpp @@ -0,0 +1,17 @@ +#include "file_descriptor.hpp" + +#include + +FileDescriptor::FileDescriptor(int value) : value_(value) {} + +FileDescriptor::~FileDescriptor() noexcept +{ + if (value_ >= 0) + close(value_); +} + +int +FileDescriptor::get() const +{ + return value_; +} diff --git a/deploy/snapshot/pagebroker/file_descriptor.hpp b/deploy/snapshot/pagebroker/file_descriptor.hpp new file mode 100644 index 000000000000..b178cf540739 --- /dev/null +++ b/deploy/snapshot/pagebroker/file_descriptor.hpp @@ -0,0 +1,15 @@ +#pragma once + +class FileDescriptor { + public: + explicit FileDescriptor(int value); + ~FileDescriptor() noexcept; + + FileDescriptor(const FileDescriptor&) = delete; + FileDescriptor& operator=(const FileDescriptor&) = delete; + + int get() const; + + private: + int value_; +}; From 9bd84414b8c11bde5578585c2f0f59f018fa7cbf Mon Sep 17 00:00:00 2001 From: Dan Feigin Date: Thu, 13 Aug 2026 14:57:23 +0300 Subject: [PATCH 2/2] chore(snapshot): satisfy PageBroker style checks --- deploy/snapshot/pagebroker/Dockerfile | 3 +++ .../snapshot/pagebroker/checkpoint_transaction_descriptor.hpp | 3 ++- deploy/snapshot/pagebroker/daemon.cpp | 3 ++- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/deploy/snapshot/pagebroker/Dockerfile b/deploy/snapshot/pagebroker/Dockerfile index ceddb2047c66..5ed43faafc91 100644 --- a/deploy/snapshot/pagebroker/Dockerfile +++ b/deploy/snapshot/pagebroker/Dockerfile @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + FROM ubuntu:24.04 AS build RUN apt-get update && apt-get install -y --no-install-recommends \ diff --git a/deploy/snapshot/pagebroker/checkpoint_transaction_descriptor.hpp b/deploy/snapshot/pagebroker/checkpoint_transaction_descriptor.hpp index 7c7bf687db40..ec63fc2d4e67 100644 --- a/deploy/snapshot/pagebroker/checkpoint_transaction_descriptor.hpp +++ b/deploy/snapshot/pagebroker/checkpoint_transaction_descriptor.hpp @@ -6,7 +6,8 @@ namespace snapshot::pagebroker { class CheckpointTransactionDescriptor { public: - CheckpointTransactionDescriptor(Path staging_directory, StorageBackend destination_storage, TransferEngineType engine_type); + CheckpointTransactionDescriptor( + Path staging_directory, StorageBackend destination_storage, TransferEngineType engine_type); const Path& staging_directory() const; const StorageBackend& destination_storage() const; diff --git a/deploy/snapshot/pagebroker/daemon.cpp b/deploy/snapshot/pagebroker/daemon.cpp index fd8357397349..8bc69b0f2200 100644 --- a/deploy/snapshot/pagebroker/daemon.cpp +++ b/deploy/snapshot/pagebroker/daemon.cpp @@ -128,7 +128,8 @@ main(int argc, char** argv) 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) { + if (bind(listener.get(), reinterpret_cast(&address), sizeof(address)) < 0 || + listen(listener.get(), 16) < 0) { std::cerr << "listen: " << std::strerror(errno) << '\n'; return 1; }