[libFuzzer] implement a better queue for the fork mode. Add an internal flag -stop_file to allow graceful shutdown of fuzzing. Enhance the logging in the fork mode

llvm-svn: 363470
diff --git a/compiler-rt/lib/fuzzer/FuzzerDriver.cpp b/compiler-rt/lib/fuzzer/FuzzerDriver.cpp
index 5458d6c..54c7ff0 100644
--- a/compiler-rt/lib/fuzzer/FuzzerDriver.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerDriver.cpp
@@ -709,6 +709,8 @@
   if (Flags.collect_data_flow)
     Options.CollectDataFlow = Flags.collect_data_flow;
   Options.LazyCounters = Flags.lazy_counters;
+  if (Flags.stop_file)
+    Options.StopFile = Flags.stop_file;
 
   unsigned Seed = Flags.seed;
   // Initialize Seed.
diff --git a/compiler-rt/lib/fuzzer/FuzzerFlags.def b/compiler-rt/lib/fuzzer/FuzzerFlags.def
index 6f0d5ad..a11cfe4 100644
--- a/compiler-rt/lib/fuzzer/FuzzerFlags.def
+++ b/compiler-rt/lib/fuzzer/FuzzerFlags.def
@@ -50,6 +50,7 @@
 FUZZER_FLAG_INT(merge, 0, "If 1, the 2-nd, 3-rd, etc corpora will be "
   "merged into the 1-st corpus. Only interesting units will be taken. "
   "This flag can be used to minimize a corpus.")
+FUZZER_FLAG_STRING(stop_file, "Stop fuzzing ASAP if this file exists")
 FUZZER_FLAG_STRING(merge_inner, "internal flag")
 FUZZER_FLAG_STRING(merge_control_file,
                    "Specify a control file used for the merge process. "
diff --git a/compiler-rt/lib/fuzzer/FuzzerFork.cpp b/compiler-rt/lib/fuzzer/FuzzerFork.cpp
index 5c4855f..95ed365 100644
--- a/compiler-rt/lib/fuzzer/FuzzerFork.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerFork.cpp
@@ -19,6 +19,7 @@
 
 #include <atomic>
 #include <chrono>
+#include <condition_variable>
 #include <fstream>
 #include <memory>
 #include <mutex>
@@ -68,6 +69,9 @@
   std::string LogPath;
   std::string SeedListPath;
   std::string CFPath;
+  size_t      JobId;
+
+  int         DftTimeInSeconds = 0;
 
   // Fuzzing Outputs.
   int ExitCode;
@@ -102,6 +106,8 @@
 
   size_t NumRuns = 0;
 
+  std::string StopFile() { return DirPlusFile(TempDir, "STOP"); }
+
   size_t secondsSinceProcessStartUp() const {
     return std::chrono::duration_cast<std::chrono::seconds>(
                std::chrono::system_clock::now() - ProcessStartTime)
@@ -119,6 +125,7 @@
     Cmd.addFlag("print_final_stats", "1");
     Cmd.addFlag("print_funcs", "0");  // no need to spend time symbolizing.
     Cmd.addFlag("max_total_time", std::to_string(std::min((size_t)300, JobId)));
+    Cmd.addFlag("stop_file", StopFile());
     if (!DataFlowBinary.empty()) {
       Cmd.addFlag("data_flow_trace", DFTDir);
       if (!Cmd.hasFlag("focus_function"))
@@ -128,11 +135,14 @@
     std::string Seeds;
     if (size_t CorpusSubsetSize =
             std::min(Files.size(), (size_t)sqrt(Files.size() + 2))) {
+      auto Time1 = std::chrono::system_clock::now();
       for (size_t i = 0; i < CorpusSubsetSize; i++) {
         auto &SF = Files[Rand->SkewTowardsLast(Files.size())];
         Seeds += (Seeds.empty() ? "" : ",") + SF;
         CollectDFT(SF);
       }
+      auto Time2 = std::chrono::system_clock::now();
+      Job->DftTimeInSeconds = duration_cast<seconds>(Time2 - Time1).count();
     }
     if (!Seeds.empty()) {
       Job->SeedListPath =
@@ -144,6 +154,7 @@
     Job->CorpusDir = DirPlusFile(TempDir, "C" + std::to_string(JobId));
     Job->FeaturesDir = DirPlusFile(TempDir, "F" + std::to_string(JobId));
     Job->CFPath = DirPlusFile(TempDir, std::to_string(JobId) + ".merge");
+    Job->JobId = JobId;
 
 
     Cmd.addArgument(Job->CorpusDir);
@@ -189,6 +200,13 @@
         }
       }
     }
+    // if (!FilesToAdd.empty() || Job->ExitCode != 0)
+    Printf("#%zd: cov: %zd ft: %zd corp: %zd exec/s %zd "
+           "oom/timeout/crash: %zd/%zd/%zd time: %zds job: %zd dft_time: %d\n",
+           NumRuns, Cov.size(), Features.size(), Files.size(),
+           Stats.average_exec_per_sec, NumOOMs, NumTimeouts, NumCrashes,
+           secondsSinceProcessStartUp(), Job->JobId, Job->DftTimeInSeconds);
+
     if (MergeCandidates.empty()) return;
 
     Vector<std::string> FilesToAdd;
@@ -209,12 +227,6 @@
           PrintPC("  NEW_FUNC: %p %F %L\n", "",
                   TPC.GetNextInstructionPc(TE->PC));
 
-    if (!FilesToAdd.empty() || Job->ExitCode != 0)
-      Printf("#%zd: cov: %zd ft: %zd corp: %zd exec/s %zd "
-             "oom/timeout/crash: %zd/%zd/%zd time: %zds\n", NumRuns,
-             Cov.size(), Features.size(), Files.size(),
-             Stats.average_exec_per_sec,
-             NumOOMs, NumTimeouts, NumCrashes, secondsSinceProcessStartUp());
   }
 
 
@@ -239,28 +251,29 @@
 struct JobQueue {
   std::queue<FuzzJob *> Qu;
   std::mutex Mu;
+  std::condition_variable Cv;
 
   void Push(FuzzJob *Job) {
-    std::lock_guard<std::mutex> Lock(Mu);
-    Qu.push(Job);
+    {
+      std::lock_guard<std::mutex> Lock(Mu);
+      Qu.push(Job);
+    }
+    Cv.notify_one();
   }
   FuzzJob *Pop() {
-    std::lock_guard<std::mutex> Lock(Mu);
-    if (Qu.empty()) return nullptr;
+    std::unique_lock<std::mutex> Lk(Mu);
+    // std::lock_guard<std::mutex> Lock(Mu);
+    Cv.wait(Lk, [&]{return !Qu.empty();});
+    assert(!Qu.empty());
     auto Job = Qu.front();
     Qu.pop();
     return Job;
   }
 };
 
-void WorkerThread(std::atomic<bool> *Stop, JobQueue *FuzzQ, JobQueue *MergeQ) {
-  while (!Stop->load()) {
-    auto Job = FuzzQ->Pop();
+void WorkerThread(JobQueue *FuzzQ, JobQueue *MergeQ) {
+  while (auto Job = FuzzQ->Pop()) {
     // Printf("WorkerThread: job %p\n", Job);
-    if (!Job) {
-      SleepSeconds(1);
-      continue;
-    }
     Job->ExitCode = ExecuteCommand(Job->Cmd);
     MergeQ->Push(Job);
   }
@@ -307,27 +320,29 @@
   int ExitCode = 0;
 
   JobQueue FuzzQ, MergeQ;
-  std::atomic<bool> Stop(false);
+
+  auto StopJobs = [&]() {
+    for (int i = 0; i < NumJobs; i++)
+      FuzzQ.Push(nullptr);
+    MergeQ.Push(nullptr);
+    WriteToFile(Unit({1}), Env.StopFile());
+  };
 
   size_t JobId = 1;
   Vector<std::thread> Threads;
   for (int t = 0; t < NumJobs; t++) {
-    Threads.push_back(std::thread(WorkerThread, &Stop, &FuzzQ, &MergeQ));
+    Threads.push_back(std::thread(WorkerThread, &FuzzQ, &MergeQ));
     FuzzQ.Push(Env.CreateNewJob(JobId++));
   }
 
   while (true) {
     std::unique_ptr<FuzzJob> Job(MergeQ.Pop());
-    if (!Job) {
-      if (Stop)
-        break;
-      SleepSeconds(1);
-      continue;
-    }
+    if (!Job)
+      break;
     ExitCode = Job->ExitCode;
     if (ExitCode == Options.InterruptExitCode) {
       Printf("==%lu== libFuzzer: a child was interrupted; exiting\n", GetPid());
-      Stop = true;
+      StopJobs();
       break;
     }
     Fuzzer::MaybeExitGracefully();
@@ -352,7 +367,8 @@
         // And exit if we don't ignore this crash.
         Printf("INFO: log from the inner process:\n%s",
                FileToString(Job->LogPath).c_str());
-        Stop = true;
+        StopJobs();
+        break;
       }
     }
 
@@ -360,22 +376,22 @@
     // This is not precise, since other threads are still running
     // and we will wait while joining them.
     // We also don't stop instantly: other jobs need to finish.
-    if (Options.MaxTotalTimeSec > 0 && !Stop &&
+    if (Options.MaxTotalTimeSec > 0 &&
         Env.secondsSinceProcessStartUp() >= (size_t)Options.MaxTotalTimeSec) {
       Printf("INFO: fuzzed for %zd seconds, wrapping up soon\n",
              Env.secondsSinceProcessStartUp());
-      Stop = true;
+      StopJobs();
+      break;
     }
-    if (!Stop && Env.NumRuns >= Options.MaxNumberOfRuns) {
+    if (Env.NumRuns >= Options.MaxNumberOfRuns) {
       Printf("INFO: fuzzed for %zd iterations, wrapping up soon\n",
              Env.NumRuns);
-      Stop = true;
+      StopJobs();
+      break;
     }
 
-    if (!Stop)
-      FuzzQ.Push(Env.CreateNewJob(JobId++));
+    FuzzQ.Push(Env.CreateNewJob(JobId++));
   }
-  Stop = true;
 
   for (auto &T : Threads)
     T.join();
diff --git a/compiler-rt/lib/fuzzer/FuzzerLoop.cpp b/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
index 7081daa..f773f9a 100644
--- a/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
@@ -801,6 +801,9 @@
 
   while (true) {
     auto Now = system_clock::now();
+    if (!Options.StopFile.empty() &&
+        !FileToVector(Options.StopFile, 1, false).empty())
+      break;
     if (duration_cast<seconds>(Now - LastCorpusReload).count() >=
         Options.ReloadIntervalSec) {
       RereadOutputCorpus(MaxInputLen);
diff --git a/compiler-rt/lib/fuzzer/FuzzerOptions.h b/compiler-rt/lib/fuzzer/FuzzerOptions.h
index 687f2ff..ad3df01 100644
--- a/compiler-rt/lib/fuzzer/FuzzerOptions.h
+++ b/compiler-rt/lib/fuzzer/FuzzerOptions.h
@@ -53,6 +53,7 @@
   std::string DataFlowTrace;
   std::string CollectDataFlow;
   std::string FeaturesDir;
+  std::string StopFile;
   bool SaveArtifacts = true;
   bool PrintNEW = true; // Print a status line when new units are found;
   bool PrintNewCovPcs = false;
diff --git a/compiler-rt/test/fuzzer/only-some-bytes-fork.test b/compiler-rt/test/fuzzer/only-some-bytes-fork.test
index 4eb0482..f224305 100644
--- a/compiler-rt/test/fuzzer/only-some-bytes-fork.test
+++ b/compiler-rt/test/fuzzer/only-some-bytes-fork.test
@@ -8,5 +8,5 @@
 
 # Test that the fork mode can collect and use the DFT
 RUN: rm -rf %t && mkdir %t
-RUN: not %t-Fuzz -collect_data_flow=%t-DFT -use_value_profile=1 -runs=100000000 -fork=1 2> %t/log
+RUN: not %t-Fuzz -collect_data_flow=%t-DFT -use_value_profile=1 -runs=100000000 -fork=20 2> %t/log
 RUN: grep BINGO %t/log