feat(tvix): Support systemd socket activation r/1755

When the nix daemon starts up, first check (using sd_listen_fds) whether
we have been systemd socket-activated. If so, instead of passing the nix
daemon socket path to grpc, start a manual accept(2) loop, passing the
client file descriptors to grpc via AddInsecureChannelFromFd. There's an
open grpc issue at https://github.com/grpc/grpc/issues/19133 for
building support into grpc to do this automatically, but as of right now
this appears to be the only way to make this happen.

Making this happen, by the way, was a bit of a journey - at one point I
attempted to ServerBuilder's experimental AddExternalConnectionAcceptor
API, and that didn't work either - it appears that the final missing
piece to getting this working was explicitly fcntl(2)ing the client file
descriptors to set O_NONBLOCK before passing them into gRPC. With that
set, this all works inside of the test vm.

Fixes: b/56
Change-Id: I5d2ab2b5b02eb570249b30a9674e115c61b0ab0e
Reviewed-on: https://cl.tvl.fyi/c/depot/+/1882
Reviewed-by: kanepyork <rikingcoding@gmail.com>
Reviewed-by: tazjin <mail@tazj.in>
Tested-by: BuildkiteCI

2 files changed, 84 insertions, 14 deletions

diff --git a/third_party/nix/src/nix-daemon/CMakeLists.txt b/third_party/nix/src/nix-daemon/CMakeLists.txt
index 27a4a1254e77..63125a9b26b2 100644
--- a/third_party/nix/src/nix-daemon/CMakeLists.txt
+++ b/third_party/nix/src/nix-daemon/CMakeLists.txt
@@ -8,6 +8,8 @@ add_executable(nix-daemon)
 include_directories(${PROJECT_BINARY_DIR}) # for config.h
 set_property(TARGET nix-daemon PROPERTY CXX_STANDARD 17)
 
+pkg_check_modules(systemd REQUIRED)
+
 target_sources(nix-daemon
   PRIVATE
     nix-daemon-proto.hh
@@ -21,6 +23,7 @@ target_link_libraries(nix-daemon
   nixmain
   absl::flags
   absl::flags_parse
+  systemd
 )
 
 install(TARGETS nix-daemon DESTINATION bin)
diff --git a/third_party/nix/src/nix-daemon/nix-daemon.cc b/third_party/nix/src/nix-daemon/nix-daemon.cc
index f4128fa4b252..fd7a553e209e 100644
--- a/third_party/nix/src/nix-daemon/nix-daemon.cc
+++ b/third_party/nix/src/nix-daemon/nix-daemon.cc
@@ -8,9 +8,11 @@
 #include <glog/logging.h>
 #include <grpcpp/security/server_credentials.h>
 #include <grpcpp/server.h>
-#include <grpcpp/server_builder_impl.h>
+#include <grpcpp/server_builder.h>
+#include <grpcpp/server_posix.h>
 #include <sys/socket.h>
 #include <sys/un.h>
+#include <systemd/sd-daemon.h>
 
 #include "libmain/shared.hh"  // TODO(tazjin): can this be removed?
 #include "libstore/globals.hh"
@@ -27,6 +29,8 @@ namespace nix::daemon {
 using grpc::Server;
 using grpc_impl::ServerBuilder;
 
+namespace {
+
 // TODO(grfn): There has to be a better way to do this - this was ported
 // verbatim from the old daemon implementation without much critical evaluation.
 static int ForwardToSocket(nix::Path socket_path) {
@@ -86,35 +90,98 @@ static int ForwardToSocket(nix::Path socket_path) {
   }
 }
 
+void SetNonBlocking(int fd) {
+  int flags = fcntl(fd, F_GETFL);  // NOLINT
+  PCHECK(flags != 0) << "Error getting socket flags";
+  PCHECK(fcntl(  // NOLINT
+             fd, F_SETFL, flags | O_NONBLOCK) == 0)
+      << "Could not set socket flags";
+}
+
+}  // namespace
+
 int RunServer() {
   Store::Params params;
   params["path-info-cache-size"] = "0";
   auto store = openStore(settings.storeUri, params);
   auto worker = NewWorkerService(*store);
-
-  std::filesystem::path socket_path(settings.nixDaemonSocketFile);
-  std::filesystem::create_directories(socket_path.parent_path());
-  auto socket_addr = absl::StrFormat("unix://%s", socket_path);
-
   ServerBuilder builder;
-  builder.AddListeningPort(socket_addr, grpc::InsecureServerCredentials());
   builder.RegisterService(worker);
 
+  auto n_fds = sd_listen_fds(0);
+
+  if (n_fds > 1) {
+    LOG(FATAL) << "Too many file descriptors (" << n_fds
+               << ") received from systemd socket activation";
+  }
+
+  std::filesystem::path socket_path;
+
+  if (n_fds == 0) {
+    socket_path = settings.nixDaemonSocketFile;
+    std::filesystem::create_directories(socket_path.parent_path());
+    auto socket_addr = absl::StrFormat("unix://%s", socket_path);
+    builder.AddListeningPort(socket_addr, grpc::InsecureServerCredentials());
+  }
+
   std::unique_ptr<Server> server(builder.BuildAndStart());
-  if (server) {
-    LOG(INFO) << "Nix daemon listening at " << socket_addr;
-    server->Wait();
-    return 0;
-  } else {
+
+  if (!server) {
+    LOG(FATAL) << "Error building server";
     return 1;
   }
+
+  // We have been systemd socket-activated - instead of asking grpc to make the
+  // socket path for us, start our own accept loop and pass file descriptors to
+  // grpc.
+  //
+  // This approach was *somewhat* adapted from
+  // https://gist.github.com/yorickvP/8d523a4df2b10c5812fa7789e82b7c1b - at some
+  // point we'd like gRPC to do it for us, though - see
+  // https://github.com/grpc/grpc/issues/19133
+  if (n_fds == 1) {
+    int socket_fd = SD_LISTEN_FDS_START;
+    // Only used for logging
+    socket_path = readLink(absl::StrFormat("/proc/self/fd/%d", socket_fd));
+
+    PCHECK(sd_notify(0, "READY=1") == 0) << "Error notifying systemd";
+    for (;;) {
+      try {
+        struct sockaddr_un remote_addr {};
+        socklen_t remote_addr_len = sizeof(remote_addr);
+        int remote_fd =
+            accept(socket_fd,
+                   reinterpret_cast<struct sockaddr*>(&remote_addr),  // NOLINT
+                   &remote_addr_len);
+        checkInterrupt();
+        if (!remote_fd) {
+          if (errno == EINTR) {
+            continue;
+          }
+          PCHECK(false) << "error accepting connection";
+        }
+
+        LOG(INFO) << "Accepted remote connection on fd " << remote_fd;
+        SetNonBlocking(remote_fd);
+        grpc::AddInsecureChannelFromFd(server.get(), remote_fd);
+      } catch (Interrupted& e) {
+        return -1;
+      } catch (Error& e) {
+        LOG(ERROR) << "error processing connection: " << e.msg();
+      }
+    }
+  }
+
+  LOG(INFO) << "Nix daemon listening at " << socket_path;
+  server->Wait();
+  return 0;
 }
 
 }  // namespace nix::daemon
 
-int main(int argc, char** argv) {
+int main(int argc, char** argv) {  // NOLINT
   FLAGS_logtostderr = true;
-  google::InitGoogleLogging(argv[0]);
+  google::InitGoogleLogging(argv[0]);  // NOLINT
 
   absl::SetFlagsUsageConfig({.version_string = [] { return nix::nixVersion; }});
   absl::ParseCommandLine(argc, argv);