From 7db734afad313b8669623138130797fad797daa7 Mon Sep 17 00:00:00 2001
From: Griffin Smith <grfn@gws.fyi>
Date: Wed, 5 Aug 2020 22:46:21 -0400
Subject: feat(tvix): Implement AddToStoreNar

Implement both the client and server sides of AddToStoreNar, using a
templated generalization of the sources and sinks we were using for
AddToStore on both ends.

Change-Id: I73d0ed34118c711b125851dff99a7518ced4af35
Reviewed-on: https://cl.tvl.fyi/c/depot/+/1686
Tested-by: BuildkiteCI
Reviewed-by: kanepyork <rikingcoding@gmail.com>
---
 third_party/nix/src/libstore/rpc-store.cc | 60 +++++++++++++++++++++----------
 third_party/nix/src/libstore/rpc-store.hh |  6 ----
 2 files changed, 42 insertions(+), 24 deletions(-)

(limited to 'third_party/nix/src/libstore')

diff --git a/third_party/nix/src/libstore/rpc-store.cc b/third_party/nix/src/libstore/rpc-store.cc
index 751eb449006d..28fc1c9eec7a 100644
--- a/third_party/nix/src/libstore/rpc-store.cc
+++ b/third_party/nix/src/libstore/rpc-store.cc
@@ -56,29 +56,32 @@ T FillFrom(const U& src) {
   return result;
 }
 
-class AddToStorePathWriterSink : public BufferedSink {
+template <typename Request>
+class RPCSink : public BufferedSink {
  public:
-  explicit AddToStorePathWriterSink(
-      std::unique_ptr<
-          grpc_impl::ClientWriter<class nix::proto::AddToStoreRequest>>&&
-          writer)
+  using Writer = grpc::ClientWriter<Request>;
+  explicit RPCSink(std::unique_ptr<Writer>&& writer)
       : writer_(std::move(writer)), good_(true) {}
 
   bool good() override { return good_; }
 
   void write(const unsigned char* data, size_t len) override {
-    proto::AddToStoreRequest req;
+    Request req;
     req.set_data(data, len);
     if (!writer_->Write(req)) {
       good_ = false;
     }
   }
 
-  grpc::Status Finish() { return writer_->Finish(); }
+  ~RPCSink() override { flush(); }
+
+  grpc::Status Finish() {
+    flush();
+    return writer_->Finish();
+  }
 
  private:
-  std::unique_ptr<grpc_impl::ClientWriter<class nix::proto::AddToStoreRequest>>
-      writer_;
+  std::unique_ptr<Writer> writer_;
   bool good_;
 };
 
@@ -319,14 +322,33 @@ void RpcStore::querySubstitutablePathInfos(const PathSet& paths,
 void RpcStore::addToStore(const ValidPathInfo& info, Source& narSource,
                           RepairFlag repair, CheckSigsFlag checkSigs,
                           std::shared_ptr<FSAccessor> accessor) {
-  throw Unsupported(absl::StrCat("Not implemented ", __func__));
-}
+  ClientContext ctx;
+  google::protobuf::Empty response;
+  auto writer = stub_->AddToStoreNar(&ctx, &response);
+
+  proto::AddToStoreNarRequest path_info_req;
+  path_info_req.mutable_path_info()->mutable_path()->set_path(info.path);
+  path_info_req.mutable_path_info()->mutable_deriver()->set_path(info.deriver);
+  path_info_req.mutable_path_info()->set_nar_hash(
+      info.narHash.to_string(Base16, false));
+  for (const auto& ref : info.references) {
+    path_info_req.mutable_path_info()->add_references(ref);
+  }
+  *path_info_req.mutable_path_info()->mutable_registration_time() =
+      TimeUtil::TimeTToTimestamp(info.registrationTime);
+  path_info_req.mutable_path_info()->set_nar_size(info.narSize);
+  path_info_req.mutable_path_info()->set_ultimate(info.ultimate);
+  for (const auto& sig : info.sigs) {
+    path_info_req.mutable_path_info()->add_sigs(sig);
+  }
+  path_info_req.mutable_path_info()->set_ca(info.ca);
+  path_info_req.mutable_path_info()->set_repair(repair);
+  path_info_req.mutable_path_info()->set_check_sigs(checkSigs);
 
-void RpcStore::addToStore(const ValidPathInfo& info,
-                          const ref<std::string>& nar, RepairFlag repair,
-                          CheckSigsFlag checkSigs,
-                          std::shared_ptr<FSAccessor> accessor) {
-  throw Unsupported(absl::StrCat("Not implemented ", __func__));
+  writer->Write(path_info_req);
+
+  RPCSink sink(std::move(writer));
+  copyNAR(narSource, sink);
 }
 
 Path RpcStore::addToStore(const std::string& name, const Path& srcPath,
@@ -349,7 +371,7 @@ Path RpcStore::addToStore(const std::string& name, const Path& srcPath,
   metadata_req.mutable_meta()->set_hash_type(HashTypeToProto(hashAlgo));
   writer->Write(metadata_req);
 
-  AddToStorePathWriterSink sink(std::move(writer));
+  RPCSink sink(std::move(writer));
   dumpPath(std::filesystem::absolute(srcPath), sink);
   sink.flush();
   SuccessOrThrow(sink.Finish());
@@ -402,7 +424,9 @@ void RpcStore::ensurePath(const Path& path) {
 }
 
 void RpcStore::addTempRoot(const Path& path) {
-  throw Unsupported(absl::StrCat("Not implemented ", __func__));
+  ClientContext ctx;
+  google::protobuf::Empty response;
+  SuccessOrThrow(stub_->AddTempRoot(&ctx, StorePath(path), &response));
 }
 
 void RpcStore::addIndirectRoot(const Path& path) {
diff --git a/third_party/nix/src/libstore/rpc-store.hh b/third_party/nix/src/libstore/rpc-store.hh
index 0ef24e22b751..dd421caf44c4 100644
--- a/third_party/nix/src/libstore/rpc-store.hh
+++ b/third_party/nix/src/libstore/rpc-store.hh
@@ -56,12 +56,6 @@ class RpcStore : public LocalFSStore, public virtual Store {
                           CheckSigsFlag checkSigs = CheckSigs,
                           std::shared_ptr<FSAccessor> accessor = 0) override;
 
-  virtual void addToStore(const ValidPathInfo& info,
-                          const ref<std::string>& nar,
-                          RepairFlag repair = NoRepair,
-                          CheckSigsFlag checkSigs = CheckSigs,
-                          std::shared_ptr<FSAccessor> accessor = 0) override;
-
   virtual Path addToStore(const std::string& name, const Path& srcPath,
                           bool recursive = true, HashType hashAlgo = htSHA256,
                           PathFilter& filter = defaultPathFilter,
-- 
cgit 1.4.1