feat: Add --recompress-threads option

3 files changed, 20 insertions, 4 deletions

diff --git a/src/core/mainoptions.cpp b/src/core/mainoptions.cpp
index 2283162f..f52f574e 100644
--- a/src/core/mainoptions.cpp
+++ b/src/core/mainoptions.cpp
@@ -48,6 +48,7 @@
 #include <algorithm>
 #include <optional>
 #include <string>
+#include <thread>
 
 #ifdef HAVE_UNISTD_H
 #  include <unistd.h>
@@ -107,6 +108,9 @@ Common options:
     -X, --recompress LEVEL     recompress the cache to level LEVEL (integer or
                                "uncompressed") using the Zstandard algorithm;
                                see "Cache compression" in the manual for details
+        --recompress-threads THREADS
+                               use up to THREADS threads when recompressing the
+                               cache; default: number of CPUs
     -o, --set-config KEY=VAL   set configuration item KEY to value VAL
     -x, --show-compression     show compression statistics
     -p, --show-config          show current configuration options in
@@ -329,6 +333,7 @@ enum {
   HASH_FILE,
   INSPECT,
   PRINT_STATS,
+  RECOMPRESS_THREADS,
   SHOW_LOG_STATS,
   TRIM_DIR,
   TRIM_MAX_SIZE,
@@ -356,6 +361,7 @@ const option long_options[] = {
   {"max-size", required_argument, nullptr, 'M'},
   {"print-stats", no_argument, nullptr, PRINT_STATS},
   {"recompress", required_argument, nullptr, 'X'},
+  {"recompress-threads", required_argument, nullptr, RECOMPRESS_THREADS},
   {"set-config", required_argument, nullptr, 'o'},
   {"show-compression", no_argument, nullptr, 'x'},
   {"show-config", no_argument, nullptr, 'p'},
@@ -378,6 +384,7 @@ process_main_options(int argc, const char* const* argv)
   uint8_t verbosity = 0;
   std::optional<std::string> evict_namespace;
   std::optional<uint64_t> evict_max_age;
+  uint32_t recompress_threads = std::thread::hardware_concurrency();
 
   // First pass: Handle non-command options that affect command options.
   while ((c = getopt_long(argc,
@@ -397,6 +404,11 @@ process_main_options(int argc, const char* const* argv)
       Util::setenv("CCACHE_CONFIGPATH", arg);
       break;
 
+    case RECOMPRESS_THREADS:
+      recompress_threads = util::value_or_throw<Error>(util::parse_unsigned(
+        arg, 1, std::numeric_limits<uint32_t>::max(), "threads"));
+      break;
+
     case TRIM_MAX_SIZE:
       trim_max_size = Util::parse_size(arg);
       break;
@@ -432,6 +444,7 @@ process_main_options(int argc, const char* const* argv)
     switch (c) {
     case CONFIG_PATH:
     case 'd': // --dir
+    case RECOMPRESS_THREADS:
     case TRIM_MAX_SIZE:
     case TRIM_METHOD:
     case 'v': // --verbose
@@ -650,7 +663,9 @@ process_main_options(int argc, const char* const* argv)
 
       ProgressBar progress_bar("Recompressing...");
       storage::local::LocalStorage(config).recompress(
-        wanted_level, [&](double progress) { progress_bar.update(progress); });
+        wanted_level, recompress_threads, [&](double progress) {
+          progress_bar.update(progress);
+        });
       break;
     }
 
diff --git a/src/storage/local/LocalStorage.hpp b/src/storage/local/LocalStorage.hpp
index 6251b3e4..54a8e9fe 100644
--- a/src/storage/local/LocalStorage.hpp
+++ b/src/storage/local/LocalStorage.hpp
@@ -104,6 +104,7 @@ public:
   get_compression_statistics(const ProgressReceiver& progress_receiver) const;
 
   void recompress(std::optional<int8_t> level,
+                  uint32_t threads,
                   const ProgressReceiver& progress_receiver);
 
 private:
diff --git a/src/storage/local/LocalStorage_compress.cpp b/src/storage/local/LocalStorage_compress.cpp
index d983b5ca..19c7cd25 100644
--- a/src/storage/local/LocalStorage_compress.cpp
+++ b/src/storage/local/LocalStorage_compress.cpp
@@ -44,7 +44,6 @@
 
 #include <memory>
 #include <string>
-#include <thread>
 
 namespace storage::local {
 
@@ -192,10 +191,11 @@ LocalStorage::get_compression_statistics(
 
 void
 LocalStorage::recompress(const std::optional<int8_t> level,
+                         const uint32_t threads,
                          const ProgressReceiver& progress_receiver)
 {
-  const size_t threads = std::thread::hardware_concurrency();
-  const size_t read_ahead = 2 * threads;
+  const size_t read_ahead =
+    std::max(static_cast<size_t>(10), 2 * static_cast<size_t>(threads));
   ThreadPool thread_pool(threads, read_ahead);
   RecompressionStatistics statistics;