reduce life time of Source objects and fix progress bar

NixOS · Jan 19, 2025 · da16041 · da16041
1 parent 765f0bc
commit da16041
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 15 deletions.
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
@@ -539,16 +539,31 @@ void RemoteStore::addMultipleToStore(
     RepairFlag repair,
     CheckSigsFlag checkSigs)
 {
-    auto source = sinkToSource([&](Sink & sink) {
-        sink << pathsToCopy.size();
-        for (auto & [pathInfo, pathSource] : pathsToCopy) {
-            WorkerProto::Serialise<ValidPathInfo>::write(*this,
-                 WorkerProto::WriteConn {
-                     .to = sink,
-                     .version = 16,
-                 },
-                 pathInfo);
-            pathSource->drainInto(sink);
+    // `addMultipleToStore` is single threaded
+    size_t nrDone = 0;
+    uint64_t bytesExpected = 0;
+    for (auto & [pathInfo, _] : pathsToCopy) {
+        bytesExpected += pathInfo.narSize;
+    }
+    act.setExpected(actCopyPath, bytesExpected);
+
+    auto source = sinkToSource([&, nrTotal = pathsToCopy.size()](Sink & sink) {
+        sink << nrTotal;
+        for (auto & [pathInfo, pathSource_] : pathsToCopy) {
+            act.progress(nrDone, nrTotal, size_t(1), size_t(0));
+            {
+                // make sure that the Source object is destroyed when we're done
+                auto pathSource = std::move(pathSource_);
+
+                WorkerProto::Serialise<ValidPathInfo>::write(*this,
+                    WorkerProto::WriteConn {
+                        .to = sink,
+                        .version = 16,
+                    },
+                    pathInfo);
+                pathSource->drainInto(sink);
+            }
+            nrDone++;
         }
     });
 

diff --git a/src/libstore/store-api.cc b/src/libstore/store-api.cc
@@ -242,8 +242,8 @@ void Store::addMultipleToStore(
         storePathsToAdd.insert(thingToAdd.first.path);
     }
 
-    auto showProgress = [&]() {
-        act.progress(nrDone, pathsToCopy.size(), nrRunning, nrFailed);
+    auto showProgress = [&, nrTotal = pathsToCopy.size()]() {
+        act.progress(nrDone, nrTotal, nrRunning, nrFailed);
     };
 
     ThreadPool pool;
@@ -1108,9 +1108,6 @@ std::map<StorePath, StorePath> copyPaths(
         return storePathForDst;
     };
 
-    // total is accessed by each copy, which are each handled in separate threads
-    std::atomic<uint64_t> total = 0;
-
     for (auto & missingPath : sortedMissing) {
         auto info = srcStore.queryPathInfo(missingPath);
 
@@ -1123,6 +1120,7 @@ std::map<StorePath, StorePath> copyPaths(
         auto source = sinkToSource([&](Sink & sink) {
             // We can reasonably assume that the copy will happen whenever we
             // read the path, so log something about that at that point
+            uint64_t total = 0;
             auto srcUri = srcStore.getUri();
             auto dstUri = dstStore.getUri();
             auto storePathS = srcStore.printStorePath(missingPath);