feat(3p/nix): Implement FindRoots, CollectGarbage r/1600

Implement the RPC client and server handlers for the FindRoots and
CollectGarbage RPC calls

Change-Id: Ifa5d582c6a33bd1e7661ac2fc860505ef404dad0
Reviewed-on: https://cl.tvl.fyi/c/depot/+/1656
Tested-by: BuildkiteCI
Reviewed-by: kanepyork <rikingcoding@gmail.com>

5 files changed, 117 insertions, 3 deletions

diff --git a/third_party/nix/src/libstore/rpc-store.cc b/third_party/nix/src/libstore/rpc-store.cc
index 0b57127610..751eb44900 100644
--- a/third_party/nix/src/libstore/rpc-store.cc
+++ b/third_party/nix/src/libstore/rpc-store.cc
@@ -418,11 +418,36 @@ void RpcStore::syncWithGC() {
 }
 
 Roots RpcStore::findRoots(bool censor) {
-  throw Unsupported(absl::StrCat("Not implemented ", __func__));
+  ClientContext ctx;
+  proto::FindRootsResponse response;
+  SuccessOrThrow(stub_->FindRoots(&ctx, kEmpty, &response));
+  Roots result;
+
+  for (const auto& [target, links] : response.roots()) {
+    auto link_paths = FillFrom<std::unordered_set<std::string>>(links.paths());
+    result.insert({target, link_paths});
+  }
+
+  return result;
 }
 
 void RpcStore::collectGarbage(const GCOptions& options, GCResults& results) {
-  throw Unsupported(absl::StrCat("Not implemented ", __func__));
+  ClientContext ctx;
+  proto::CollectGarbageRequest request;
+  request.set_action(options.ActionToProto());
+  for (const auto& path : options.pathsToDelete) {
+    request.add_paths_to_delete(path);
+  }
+  request.set_ignore_liveness(options.ignoreLiveness);
+  request.set_max_freed(options.maxFreed);
+
+  proto::CollectGarbageResponse response;
+  SuccessOrThrow(stub_->CollectGarbage(&ctx, request, &response));
+
+  for (const auto& path : response.deleted_paths()) {
+    results.paths.insert(path);
+  }
+  results.bytesFreed = response.bytes_freed();
 }
 
 void RpcStore::optimiseStore() {
diff --git a/third_party/nix/src/libstore/store-api.cc b/third_party/nix/src/libstore/store-api.cc
index 444ec73cc7..ab71282272 100644
--- a/third_party/nix/src/libstore/store-api.cc
+++ b/third_party/nix/src/libstore/store-api.cc
@@ -77,6 +77,35 @@ nix::proto::BuildStatus BuildResult::status_to_proto() {
   }
 }
 
+std::optional<GCOptions::GCAction> GCActionFromProto(
+    nix::proto::GCAction gc_action) {
+  switch (gc_action) {
+    case nix::proto::GCAction::ReturnLive:
+      return GCOptions::GCAction::gcReturnLive;
+    case nix::proto::GCAction::ReturnDead:
+      return GCOptions::GCAction::gcReturnDead;
+    case nix::proto::GCAction::DeleteDead:
+      return GCOptions::GCAction::gcDeleteDead;
+    case nix::proto::GCAction::DeleteSpecific:
+      return GCOptions::GCAction::gcDeleteSpecific;
+    default:
+      return {};
+  }
+}
+
+[[nodiscard]] const proto::GCAction GCOptions::ActionToProto() const {
+  switch (action) {
+    case GCOptions::GCAction::gcReturnLive:
+      return nix::proto::GCAction::ReturnLive;
+    case GCOptions::GCAction::gcReturnDead:
+      return nix::proto::GCAction::ReturnDead;
+    case GCOptions::GCAction::gcDeleteDead:
+      return nix::proto::GCAction::DeleteDead;
+    case GCOptions::GCAction::gcDeleteSpecific:
+      return nix::proto::GCAction::DeleteSpecific;
+  }
+}
+
 bool Store::isInStore(const Path& path) const {
   return isInDir(path, storeDir);
 }
diff --git a/third_party/nix/src/libstore/store-api.hh b/third_party/nix/src/libstore/store-api.hh
index 7764ae8b83..974734dc15 100644
--- a/third_party/nix/src/libstore/store-api.hh
+++ b/third_party/nix/src/libstore/store-api.hh
@@ -80,8 +80,13 @@ struct GCOptions {
 
   /* Stop after at least `maxFreed' bytes have been freed. */
   unsigned long long maxFreed{std::numeric_limits<unsigned long long>::max()};
+
+  [[nodiscard]] const proto::GCAction ActionToProto() const;
 };
 
+std::optional<GCOptions::GCAction> GCActionFromProto(
+    nix::proto::GCAction gc_action);
+
 struct GCResults {
   /* Depending on the action, the GC roots, or the paths that would
      be or have been deleted. */
diff --git a/third_party/nix/src/nix-daemon/nix-daemon-proto.cc b/third_party/nix/src/nix-daemon/nix-daemon-proto.cc
index fe580b08a2..b310cc92c4 100644
--- a/third_party/nix/src/nix-daemon/nix-daemon-proto.cc
+++ b/third_party/nix/src/nix-daemon/nix-daemon-proto.cc
@@ -252,6 +252,61 @@ class WorkerServiceImpl final : public WorkerService::Service {
         __FUNCTION__);
   }
 
+  Status FindRoots(grpc::ServerContext*, const google::protobuf::Empty*,
+                   nix::proto::FindRootsResponse* response) override {
+    return HandleExceptions(
+        [&]() -> Status {
+          auto roots = store_->findRoots(false);
+          for (const auto& [target, links] : roots) {
+            StorePaths link_paths;
+            for (const auto& link : links) {
+              link_paths.add_paths(link);
+            }
+            response->mutable_roots()->insert({target, link_paths});
+          }
+
+          return Status::OK;
+        },
+        __FUNCTION__);
+  }
+
+  Status CollectGarbage(grpc::ServerContext*,
+                        const proto::CollectGarbageRequest* request,
+                        proto::CollectGarbageResponse* response) override {
+    return HandleExceptions(
+        [&]() -> Status {
+          GCOptions options;
+          auto action = GCActionFromProto(request->action());
+          if (!action.has_value()) {
+            return Status(grpc::StatusCode::INVALID_ARGUMENT,
+                          "Invalid GC action");
+          }
+
+          options.action = action.value();
+          for (const auto& path : request->paths_to_delete()) {
+            options.pathsToDelete.insert(path);
+          }
+          options.ignoreLiveness = request->ignore_liveness();
+          options.maxFreed = request->max_freed();
+
+          if (options.ignoreLiveness) {
+            return Status(grpc::StatusCode::INVALID_ARGUMENT,
+                          "you are not allowed to ignore liveness");
+          }
+
+          GCResults results;
+          store_->collectGarbage(options, results);
+
+          for (const auto& path : results.paths) {
+            response->add_deleted_paths(path);
+          }
+          response->set_bytes_freed(results.bytesFreed);
+
+          return Status::OK;
+        },
+        __FUNCTION__);
+  }
+
   Status QuerySubstitutablePathInfos(
       grpc::ServerContext*, const StorePaths* request,
       nix::proto::SubstitutablePathInfos* response) override {
diff --git a/third_party/nix/src/proto/worker.proto b/third_party/nix/src/proto/worker.proto
index e632c6264f..abe8f622c9 100644
--- a/third_party/nix/src/proto/worker.proto
+++ b/third_party/nix/src/proto/worker.proto
@@ -203,7 +203,7 @@ message BuildPathsRequest {
 }
 
 message FindRootsResponse {
-  map<string, string> roots = 1;
+  map<string, StorePaths> roots = 1;
 }
 
 message SetOptionsRequest {