feat(3p/nix): Implement a few more RPC calls r/1529

Implement the RPC client calls for QueryPathFromHashPart,
QuerySubstitutablePaths, and QuerySubstitutablePathInfos, and the
handler for QuerySubstitutablePathInfos.

Refs: #29
Change-Id: Idf383b771f159f267d8f65367bc4af3d239e32b7
Reviewed-on: https://cl.tvl.fyi/c/depot/+/1515
Tested-by: BuildkiteCI
Reviewed-by: kanepyork <rikingcoding@gmail.com>

3 files changed, 59 insertions, 3 deletions

diff --git a/third_party/nix/src/libstore/rpc-store.cc b/third_party/nix/src/libstore/rpc-store.cc
index e2eb240f8b..d1687473c1 100644
--- a/third_party/nix/src/libstore/rpc-store.cc
+++ b/third_party/nix/src/libstore/rpc-store.cc
@@ -38,6 +38,14 @@ proto::StorePath StorePath(const Path& path) {
   return store_path;
 }
 
+proto::StorePaths StorePaths(const PathSet& paths) {
+  proto::StorePaths result;
+  for (const auto& path : paths) {
+    result.add_paths(path);
+  }
+  return result;
+}
+
 template <typename T, typename U>
 T FillFrom(const U& src) {
   T result;
@@ -143,16 +151,37 @@ StringSet RpcStore::queryDerivationOutputNames(const Path& path) {
 }
 
 Path RpcStore::queryPathFromHashPart(const std::string& hashPart) {
-  throw absl::StrCat("Not implemented ", __func__);
+  proto::StorePath path;
+  proto::HashPart proto_hash_part;
+  proto_hash_part.set_hash_part(hashPart);
+  SuccessOrThrow(stub_->QueryPathFromHashPart(&ctx, proto_hash_part, &path));
+  return path.path();
 }
 
 PathSet RpcStore::querySubstitutablePaths(const PathSet& paths) {
-  throw absl::StrCat("Not implemented ", __func__);
+  proto::StorePaths result;
+  SuccessOrThrow(
+      stub_->QuerySubstitutablePaths(&ctx, StorePaths(paths), &result));
+  return FillFrom<PathSet>(result.paths());
 }
 
 void RpcStore::querySubstitutablePathInfos(const PathSet& paths,
                                            SubstitutablePathInfos& infos) {
-  throw absl::StrCat("Not implemented ", __func__);
+  proto::SubstitutablePathInfos result;
+  SuccessOrThrow(
+      stub_->QuerySubstitutablePathInfos(&ctx, StorePaths(paths), &result));
+
+  for (const auto& path_info : result.path_infos()) {
+    auto path = path_info.path().path();
+    SubstitutablePathInfo& info(infos[path]);
+    info.deriver = path_info.deriver().path();
+    if (!info.deriver.empty()) {
+      assertStorePath(info.deriver);
+    }
+    info.references = FillFrom<PathSet>(path_info.references());
+    info.downloadSize = path_info.download_size();
+    info.narSize = path_info.nar_size();
+  }
 }
 
 void RpcStore::addToStore(const ValidPathInfo& info, Source& narSource,
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 27c694b292..58f7754386 100644
--- a/third_party/nix/src/nix-daemon/nix-daemon-proto.cc
+++ b/third_party/nix/src/nix-daemon/nix-daemon-proto.cc
@@ -17,6 +17,7 @@
 #include "libutil/archive.hh"
 #include "libutil/hash.hh"
 #include "libutil/serialise.hh"
+#include "libutil/types.hh"
 
 namespace nix::daemon {
 
@@ -170,6 +171,29 @@ class WorkerServiceImpl final : public WorkerService::Service {
     return Status::OK;
   }
 
+  Status QuerySubstitutablePathInfos(
+      grpc::ServerContext*, const StorePaths* request,
+      nix::proto::SubstitutablePathInfos* response) override {
+    SubstitutablePathInfos infos;
+    PathSet paths;
+    for (const auto& path : request->paths()) {
+      paths.insert(path);
+    }
+    store_->querySubstitutablePathInfos(paths, infos);
+    for (const auto& [path, path_info] : infos) {
+      auto proto_path_info = response->add_path_infos();
+      proto_path_info->mutable_path()->set_path(path);
+      proto_path_info->mutable_deriver()->set_path(path_info.deriver);
+      for (const auto& ref : path_info.references) {
+        proto_path_info->add_references(ref);
+      }
+      proto_path_info->set_download_size(path_info.downloadSize);
+      proto_path_info->set_nar_size(path_info.narSize);
+    }
+
+    return Status::OK;
+  }
+
   Status QueryValidDerivers(grpc::ServerContext* context,
                             const StorePath* request,
                             StorePaths* response) override {
@@ -216,6 +240,7 @@ class WorkerServiceImpl final : public WorkerService::Service {
                        PathInfo* response) override {
     auto path = request->path();
     store_->assertStorePath(path);
+    response->mutable_path()->set_path(path);
     try {
       auto info = store_->queryPathInfo(path);
       response->mutable_deriver()->set_path(info->deriver);
diff --git a/third_party/nix/src/proto/worker.proto b/third_party/nix/src/proto/worker.proto
index e2ff598b2c..e632c6264f 100644
--- a/third_party/nix/src/proto/worker.proto
+++ b/third_party/nix/src/proto/worker.proto
@@ -246,11 +246,13 @@ message CollectGarbageResponse {
 }
 
 message PathInfo {
+  StorePath path = 10;
   bool is_valid = 9;
   StorePath deriver = 1;
   bytes nar_hash = 2;
   repeated string references = 3;
   google.protobuf.Timestamp registration_time = 4;
+  uint64 download_size = 11;
   uint64 nar_size = 5;
   // Whether the path is ultimately trusted, that is, it's a derivation
   // output that was built locally.