feat: Add progress bar for indexing & merging (#402)

Author: varungandhi-src

indexer/CliOptions.h

@@ -32,6 +32,7 @@ struct CliOptions {
   std::string statsFilePath;
   std::string packageMapPath;
   bool showCompilerDiagnostics;
+  bool showProgress;
 
   size_t ipcSizeHintBytes;
   std::chrono::seconds receiveTimeout;

indexer/Driver.cc

@@ -51,6 +51,7 @@
 #include "indexer/LlvmAdapter.h"
 #include "indexer/Logging.h"
 #include "indexer/Path.h"
+#include "indexer/ProgressReporter.h"
 #include "indexer/RAII.h"
 #include "indexer/ScipExtras.h"
 #include "indexer/Statistics.h"
@@ -210,6 +211,7 @@ struct DriverOptions {
   AbsolutePath statsFilePath;
   AbsolutePath packageMapPath;
   bool showCompilerDiagnostics;
+  bool showProgress;
   DriverIpcOptions ipcOptions;
   size_t numWorkers;
   bool deterministic;
@@ -229,7 +231,8 @@ struct DriverOptions {
         projectRootPath(AbsolutePath("/"), RootKind::Project), compdbPath(),
         indexOutputPath(), statsFilePath(), packageMapPath(),
         showCompilerDiagnostics(cliOpts.showCompilerDiagnostics),
-        ipcOptions{cliOpts.ipcSizeHintBytes, cliOpts.receiveTimeout},
+        showProgress(cliOpts.showProgress), ipcOptions{cliOpts.ipcSizeHintBytes,
+                                                       cliOpts.receiveTimeout},
         numWorkers(cliOpts.numWorkers), deterministic(cliOpts.deterministic),
         preprocessorRecordHistoryFilterRegex(
             cliOpts.preprocessorRecordHistoryFilterRegex),
@@ -336,6 +339,9 @@ struct DriverOptions {
     if (this->showCompilerDiagnostics) {
       args.push_back("--show-compiler-diagnostics");
     }
+    if (!this->showProgress) {
+      args.push_back("--no-progress-report");
+    }
     if (!this->preprocessorRecordHistoryFilterRegex.empty()) {
       args.push_back(fmt::format("--preprocessor-record-history-filter={}",
                                  this->preprocessorRecordHistoryFilterRegex));
@@ -812,6 +818,13 @@ class Scheduler final {
     return this->maybeErroredJobs.size();
   }
 
+  std::string_view getTuPath(JobId jobId) const {
+    auto it = this->allJobList.find(jobId);
+    ENFORCE(it != this->allJobList.end());
+    ENFORCE(it->second.job.kind == IndexJob::Kind::SemanticAnalysis);
+    return std::string_view(it->second.job.semanticAnalysis.command.filePath);
+  }
+
 private:
   ToBeScheduledWorkerId claimIdleWorker() {
     ENFORCE(!this->idleWorkers.empty());
@@ -883,6 +896,7 @@ class Driver {
 
   /// Total number of commands in the compilation database.
   size_t compdbCommandCount = 0;
+  TusIndexedCount indexedSoFar;
   compdb::ResumableParser compdbParser;
 
 public:
@@ -1079,23 +1093,37 @@ class Driver {
     llvm::UniqueStringSaver stringSaver{allocator};
     scip::SymbolNameInterner interner{stringSaver};
     scip::IndexBuilder builder{interner};
-    // TODO: Measure how much time this is taking and parallelize if too slow.
-    for (auto &paths : this->shardPaths) {
-      scip::Index indexShard;
-      if (!readIndexShard(paths.docsAndExternals, indexShard)) {
-        continue;
-      }
-      for (auto &doc : *indexShard.mutable_documents()) {
-        bool isMultiplyIndexed = this->isMultiplyIndexedApproximate(
-            doc.relative_path(), paths.docsAndExternals.asRef(), badJobIds);
-        builder.addDocument(std::move(doc), isMultiplyIndexed);
-      }
-      // See NOTE(ref: precondition-deterministic-ext-symbol-docs); in
-      // deterministic mode, indexes should be the same, and iterated over in
-      // sorted order. So if external symbol emission in each part is
-      // deterministic, addExternalSymbol will be called in deterministic order.
-      for (auto &extSym : *indexShard.mutable_external_symbols()) {
-        builder.addExternalSymbol(std::move(extSym));
+    {
+      ProgressReporter progressReporter(this->options.showProgress,
+                                        "Merged partial index for",
+                                        this->shardPaths.size());
+      size_t count = 1;
+      for (auto &paths : this->shardPaths) {
+        scip::Index indexShard;
+        if (!readIndexShard(paths.docsAndExternals, indexShard)) {
+          continue;
+        }
+        for (auto &doc : *indexShard.mutable_documents()) {
+          bool isMultiplyIndexed = this->isMultiplyIndexedApproximate(
+              doc.relative_path(), paths.docsAndExternals.asRef(), badJobIds);
+          builder.addDocument(std::move(doc), isMultiplyIndexed);
+        }
+        // See NOTE(ref: precondition-deterministic-ext-symbol-docs); in
+        // deterministic mode, indexes should be the same, and iterated over in
+        // sorted order. So if external symbol emission in each part is
+        // deterministic, addExternalSymbol will be called in deterministic
+        // order.
+        for (auto &extSym : *indexShard.mutable_external_symbols()) {
+          builder.addExternalSymbol(std::move(extSym));
+        }
+        if (auto optFileName = paths.docsAndExternals.asRef().fileName()) {
+          if (auto optJobId = ShardPaths::tryParseJobId(optFileName.value())) {
+            progressReporter.report(
+                count,
+                this->scheduler.getTuPath(JobId::newTask(optJobId.value())));
+          }
+        }
+        count++;
       }
     }
 
@@ -1178,19 +1206,20 @@ class Driver {
 
   /// Returns the number of TUs processed
   std::pair<TusIndexedCount, size_t> runJobsTillCompletionAndShutdownWorkers() {
-    TusIndexedCount tusIndexedCount{};
+    ProgressReporter progressReporter(this->options.showProgress, "Indexed",
+                                      this->compdbCommandCount);
     this->scheduler.runJobsTillCompletionAndShutdownWorkers(
         Scheduler::RunCallbacks{
-            [this, &tusIndexedCount]() -> void {
-              tusIndexedCount += this->processOneOrMoreJobResults();
+            [this, &progressReporter]() -> void {
+              this->processOneOrMoreJobResults(progressReporter);
             },
             [this]() -> size_t { return this->refillJobs(); },
             [this](ToBeScheduledWorkerId &&workerId, JobId jobId) -> bool {
               return this->tryAssignJobToWorker(std::move(workerId), jobId);
             },
             [this](WorkerId workerId) { this->shutdownWorker(workerId); }});
     this->scheduler.waitForAllWorkers();
-    return {tusIndexedCount, this->scheduler.numErroredJobs()};
+    return {this->indexedSoFar, this->scheduler.numErroredJobs()};
   }
 
   FileGuard openCompilationDatabase() {
@@ -1241,17 +1270,17 @@ class Driver {
         });
   }
 
-  TusIndexedCount processWorkerResponse(IndexJobResponse &&response) {
+  void processWorkerResponse(IndexJobResponse &&response,
+                             const ProgressReporter &progressReporter) {
     TRACE_EVENT(tracing::indexing, "Driver::processWorkerResponse",
                 perfetto::TerminatingFlow::Global(response.jobId.traceId()));
-    TusIndexedCount tusIndexedCount{};
     auto optLatestIdleWorkerId = this->scheduler.markCompleted(
         response.workerId, response.jobId, response.result.kind);
     if (!optLatestIdleWorkerId.has_value()) {
       spdlog::debug("worker {} was terminated before job {}'s result "
                     "was processed, so can't send an emit index job",
                     response.workerId, response.jobId);
-      return tusIndexedCount;
+      return;
     }
     auto latestIdleWorkerId = *optLatestIdleWorkerId;
     switch (response.result.kind) {
@@ -1299,7 +1328,7 @@ class Driver {
               oldHandle.terminate();
               return this->spawnWorker(workerId);
             });
-        return tusIndexedCount;
+        return;
       }
       break;
     }
@@ -1310,17 +1339,21 @@ class Driver {
                                          std::move(result.statistics));
       }
       this->shardPaths.emplace_back(std::move(result.shardPaths));
-      tusIndexedCount.value += 1;
+      this->indexedSoFar.value += 1;
+      if (this->options.showProgress) {
+        auto semaJobId = JobId::newTask(response.jobId.taskId());
+        progressReporter.report(this->indexedSoFar.value,
+                                this->scheduler.getTuPath(semaJobId));
+      }
       break;
     }
     }
-    return tusIndexedCount;
+    return;
   }
 
-  TusIndexedCount processOneOrMoreJobResults() {
+  void processOneOrMoreJobResults(const ProgressReporter &progressReporter) {
     using namespace std::chrono_literals;
     auto workerTimeout = this->receiveTimeout();
-    TusIndexedCount tusIndexedCount{};
     IndexJobResponse response;
     TRACE_EVENT_BEGIN(tracing::ipc, "driver.waitForResponse");
     auto recvError =
@@ -1337,9 +1370,9 @@ class Driver {
     } else {
       spdlog::debug("received response for {} from worker {}", response.jobId,
                     response.workerId);
-      tusIndexedCount = this->processWorkerResponse(std::move(response));
+      this->processWorkerResponse(std::move(response), progressReporter);
       while (this->queues.workerToDriver.tryReceiveInstant(response)) {
-        tusIndexedCount += this->processWorkerResponse(std::move(response));
+        this->processWorkerResponse(std::move(response), progressReporter);
       }
     }
     // NOTE(def: mail-from-the-dead)
@@ -1360,7 +1393,7 @@ class Driver {
     // was actually submitted by a worker which was terminated.
     auto now = std::chrono::steady_clock::now();
     this->terminateLongRunningWorkersAndRespawn(now - workerTimeout);
-    return tusIndexedCount;
+    return;
   }
 
   // Assign a job to a specific worker. When this method is called,

indexer/IpcMessages.cc

@@ -184,7 +184,6 @@ std::optional<uint32_t> ShardPaths::tryParseJobId(std::string_view fileName) {
     return taskId;
   }
   return {};
-  // return fmt::format("job-{}-worker-{}", taskId, workerId);
 }
 
 } // namespace scip_clang

indexer/ProgressReporter.cc

@@ -0,0 +1,41 @@
+#include <cmath>
+#include <cstdint>
+#include <iostream>
+#include <ostream>
+#include <string_view>
+
+#include "spdlog/fmt/fmt.h"
+
+#include "indexer/ProgressReporter.h"
+
+namespace scip_clang {
+
+ProgressReporter::ProgressReporter(bool active, std::string_view msg,
+                                   size_t totalCount)
+    : message(msg), totalCount(totalCount), countWidth(), active(active) {
+  if (this->totalCount == 0) {
+    countWidth = 1;
+  } else {
+    countWidth = std::log10(double(this->totalCount)) + 1;
+  }
+}
+
+void ProgressReporter::report(size_t count, std::string_view extraData) const {
+  if (!this->active) {
+    return;
+  }
+  int maxExtraWidth = 256;
+  auto backspaceCount = std::max(maxExtraWidth - int(extraData.size()), 0);
+  fmt::print("\r[{1:>{0}}/{2:>{0}}] {3} {4:<{5}}{6:\b<{7}}", countWidth, count,
+             this->totalCount, this->message, extraData, maxExtraWidth, "",
+             backspaceCount);
+  std::flush(std::cout);
+}
+
+ProgressReporter::~ProgressReporter() {
+  if (this->active) {
+    fmt::print("\n");
+  }
+}
+
+} // namespace scip_clang

indexer/ProgressReporter.h

@@ -0,0 +1,25 @@
+#ifndef SCIP_CLANG_PROGRESS_REPORTER_H
+#define SCIP_CLANG_PROGRESS_REPORTER_H
+
+#include <cstdint>
+#include <string_view>
+
+namespace scip_clang {
+
+class ProgressReporter {
+  std::string_view message;
+  size_t totalCount;
+  size_t countWidth;
+  bool active;
+
+public:
+  ProgressReporter(bool active, std::string_view msg, size_t totalCount);
+
+  void report(size_t count, std::string_view extraData) const;
+
+  ~ProgressReporter();
+};
+
+} // namespace scip_clang
+
+#endif // SCIP_CLANG_PROGRESS_REPORTER_H

indexer/Worker.cc

@@ -119,6 +119,7 @@ WorkerOptions WorkerOptions::fromCliOptions(const CliOptions &cliOptions) {
                        statsFilePath,
                        cliOptions.packageMapPath,
                        cliOptions.showCompilerDiagnostics,
+                       cliOptions.showProgress,
                        cliOptions.logLevel,
                        cliOptions.deterministic,
                        cliOptions.measureStatistics,

indexer/Worker.h

@@ -58,6 +58,7 @@ struct WorkerOptions {
 
   StdPath packageMapPath;
   bool showCompilerDiagnostics;
+  bool showProgress;
 
   spdlog::level::level_enum logLevel;
   bool deterministic;

indexer/main.cc

@@ -39,6 +39,10 @@ static scip_clang::CliOptions parseArguments(int argc, char *argv[]) {
     "index-output-path",
     "Path to write the SCIP index to",
     cxxopts::value<std::string>(cliOptions.indexOutputPath)->default_value("index.scip"));
+  parser.add_options(defaultGroup)(
+    "package-map-path",
+    "Path to use for path->package mappings in JSON format.",
+    cxxopts::value<std::string>(cliOptions.packageMapPath));
   parser.add_options(defaultGroup)(
     "j,jobs",
     fmt::format(
@@ -50,9 +54,9 @@ static scip_clang::CliOptions parseArguments(int argc, char *argv[]) {
     "One of 'debug', 'info', 'warning' or 'error'",
     cxxopts::value<std::string>()->default_value("info"));
   parser.add_options(defaultGroup)(
-    "package-map-path",
-    "Path to use for path->package mappings in JSON format.",
-    cxxopts::value<std::string>(cliOptions.packageMapPath));
+    "no-progress-report",
+    "Suppress progress information reported after a TU is indexed",
+    cxxopts::value<bool>());
   parser.add_options(defaultGroup)(
     "show-compiler-diagnostics",
     "Show Clang diagnostics triggered when running semantic analysis."
@@ -172,6 +176,8 @@ static scip_clang::CliOptions parseArguments(int argc, char *argv[]) {
     std::exit(EXIT_FAILURE);
   }
 
+  cliOptions.showProgress = result.count("no-progress-report") == 0;
+
   auto level = result["log-level"].as<std::string>();
   if (level == "debug") {
     cliOptions.logLevel = spdlog::level::level_enum::debug;

test/test_main.cc

@@ -454,6 +454,7 @@ TEST_CASE("INDEX") {
         args.push_back(fmt::format("--driver-id=index-{}",
                                    test::globalCliOptions.testName));
         args.push_back("--deterministic");
+        args.push_back("--no-progress-report");
         auto packageMapPath =
             StdPath(rootInSourceDir.asStringRef()) / "package-map.json";
         if (std::filesystem::exists(packageMapPath)) {