[libFuzzer] add threads to the fork mode: now you can pass -fork=N to run N concurrent workers. Fork mode is still work-in-progress.

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@353997 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 147 insertions, 90 deletions

diff --git a/lib/fuzzer/FuzzerDriver.cpp b/lib/fuzzer/FuzzerDriver.cpp
index 00dae3fda..434c48128 100644
--- a/lib/fuzzer/FuzzerDriver.cpp
+++ b/lib/fuzzer/FuzzerDriver.cpp
@@ -725,7 +725,7 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
   }
 
   if (Flags.fork)
-    FuzzWithFork(F->GetMD().GetRand(), Options, Args, *Inputs);
+    FuzzWithFork(F->GetMD().GetRand(), Options, Args, *Inputs, Flags.fork);
 
   if (Flags.merge)
     Merge(F, Options, Args, *Inputs, Flags.merge_control_file);
diff --git a/lib/fuzzer/FuzzerFork.cpp b/lib/fuzzer/FuzzerFork.cpp
index 5126fe74c..8584e0b65 100644
--- a/lib/fuzzer/FuzzerFork.cpp
+++ b/lib/fuzzer/FuzzerFork.cpp
@@ -15,146 +15,203 @@
 #include "FuzzerSHA1.h"
 #include "FuzzerUtil.h"
 
+#include <mutex>
+#include <thread>
+#include <queue>
+
 namespace fuzzer {
 
 struct FuzzJob {
   // Inputs.
   Command Cmd;
-  Vector<std::string> Files;
   std::string CorpusDir;
   std::string LogPath;
   std::string CFPath;
-  int MaxTotalTimeSec;
 
   // Fuzzing Outputs.
   int ExitCode;
 };
 
 struct GlobalEnv {
-  const Vector<std::string> *Args;
+  Vector<std::string> Args;
+  Vector<std::string> CorpusDirs;
   std::string MainCorpusDir;
+  std::string TempDir;
   Set<uint32_t> Features;
   Vector<std::string> Files;
+  Random *Rand;
+  int Verbosity = 0;
+
+  FuzzJob *CreateNewJob(size_t JobId) {
+    Command Cmd(Args);
+    Cmd.removeFlag("fork");
+    for (auto &C : CorpusDirs) // Remove all corpora from the args.
+      Cmd.removeArgument(C);
+    Cmd.addFlag("reload", "0");  // working in an isolated dir, no reload.
+    Cmd.addFlag("max_total_time", std::to_string(std::min((size_t)300, JobId)));
+
+    auto Job = new FuzzJob;
+    std::string Seeds;
+    if (size_t CorpusSubsetSize = std::min(Files.size(), (size_t)100))
+      for (size_t i = 0; i < CorpusSubsetSize; i++)
+        Seeds += (Seeds.empty() ? "" : ",") +
+                 Files[Rand->SkewTowardsLast(Files.size())];
+    if (!Seeds.empty())
+      Cmd.addFlag("seed_inputs", Seeds);
+    Job->LogPath = DirPlusFile(TempDir, std::to_string(JobId) + ".log");
+    Job->CorpusDir = DirPlusFile(TempDir, "C" + std::to_string(JobId));
+    Job->CFPath = DirPlusFile(TempDir, std::to_string(JobId) + ".merge");
+
+
+    Cmd.addArgument(Job->CorpusDir);
+    RmDirRecursive(Job->CorpusDir);
+    MkDir(Job->CorpusDir);
+
+    Cmd.setOutputFile(Job->LogPath);
+    Cmd.combineOutAndErr();
+
+    Job->Cmd = Cmd;
+
+    if (Verbosity >= 2)
+      Printf("Job %zd/%p Created: %s\n", JobId, Job,
+             Job->Cmd.toString().c_str());
+    // Start from very short runs and gradually increase them.
+    return Job;
+  }
+
+  void RunOneMergeJob(FuzzJob *Job) {
+    Vector<SizedFile> TempFiles;
+    GetSizedFilesFromDir(Job->CorpusDir, &TempFiles);
+
+    Vector<std::string> FilesToAdd;
+    Set<uint32_t> NewFeatures;
+    CrashResistantMerge(Args, {}, TempFiles, &FilesToAdd, Features,
+                        &NewFeatures, Job->CFPath, false);
+    RemoveFile(Job->CFPath);
+    for (auto &Path : FilesToAdd) {
+      auto U = FileToVector(Path);
+      auto NewPath = DirPlusFile(MainCorpusDir, Hash(U));
+      WriteToFile(U, NewPath);
+      Files.push_back(NewPath);
+    }
+    Printf("Removing %s\n", Job->CorpusDir.c_str());
+    RmDirRecursive(Job->CorpusDir);
+    Features.insert(NewFeatures.begin(), NewFeatures.end());
+    Printf("INFO: temp_files: %zd files_added: %zd newft: %zd ft: %zd\n",
+           TempFiles.size(), FilesToAdd.size(), NewFeatures.size(),
+           Features.size());
+  }
 };
 
-void RunOneFuzzingJob(FuzzJob *Job) {
-  Command &Cmd = Job->Cmd;
-  if (!Job->Files.empty()) {
-    std::string Seeds;
-    for (const auto &File : Job->Files)
-        Seeds += (Seeds.empty() ? "" : ",") + File;
-    Cmd.addFlag("seed_inputs", Seeds);
+struct JobQueue {
+  std::queue<FuzzJob *> Qu;
+  std::mutex Mu;
+
+  void Push(FuzzJob *Job) {
+    std::lock_guard<std::mutex> Lock(Mu);
+    Qu.push(Job);
   }
-  Cmd.addFlag("max_total_time", std::to_string(Job->MaxTotalTimeSec));
-  Cmd.setOutputFile(Job->LogPath);
-  Cmd.combineOutAndErr();
-  Cmd.addArgument(Job->CorpusDir);
-  RmDirRecursive(Job->CorpusDir);
-  MkDir(Job->CorpusDir);
-  Job->ExitCode = ExecuteCommand(Cmd);
-}
+  FuzzJob *Pop() {
+    std::lock_guard<std::mutex> Lock(Mu);
+    if (Qu.empty()) return nullptr;
+    auto Job = Qu.front();
+    Qu.pop();
+    return Job;
+  }
+};
 
-void RunOneMergeJob(GlobalEnv *Env, FuzzJob *Job) {
-  Vector<SizedFile> TempFiles;
-  GetSizedFilesFromDir(Job->CorpusDir, &TempFiles);
-
-  Vector<std::string>FilesToAdd;
-  Set<uint32_t> NewFeatures;
-  CrashResistantMerge(*Env->Args, {}, TempFiles, &FilesToAdd, Env->Features,
-                      &NewFeatures, Job->CFPath, false);
-  RemoveFile(Job->CFPath);
-  for (auto &Path : FilesToAdd) {
-    auto U = FileToVector(Path);
-    auto NewPath = DirPlusFile(Env->MainCorpusDir, Hash(U));
-    WriteToFile(U, NewPath);
-    Env->Files.push_back(NewPath);
+void WorkerThread(std::atomic<bool> *Stop, JobQueue *FuzzQ, JobQueue *MergeQ) {
+  while (!*Stop) {
+    auto Job = FuzzQ->Pop();
+    // Printf("WorkerThread: job %p\n", Job);
+    if (!Job) {
+      SleepSeconds(1);
+      continue;
+    }
+    Job->ExitCode = ExecuteCommand(Job->Cmd);
+    MergeQ->Push(Job);
   }
-  RmDirRecursive(Job->CorpusDir);
-  Env->Features.insert(NewFeatures.begin(), NewFeatures.end());
-  Printf("INFO: temp_files: %zd files_added: %zd newft: %zd ft: %zd\n",
-         TempFiles.size(), FilesToAdd.size(), NewFeatures.size(),
-         Env->Features.size());
 }
 
 // This is just a skeleton of an experimental -fork=1 feature.
 void FuzzWithFork(Random &Rand, const FuzzingOptions &Options,
                   const Vector<std::string> &Args,
-                  const Vector<std::string> &CorpusDirs) {
-  Printf("INFO: -fork=1: doing fuzzing in a separate process in order to "
-         "be more resistant to crashes, timeouts, and OOMs\n");
+                  const Vector<std::string> &CorpusDirs, int NumJobs) {
+  Printf("INFO: -fork=%d: doing fuzzing in a separate process in order to "
+         "be more resistant to crashes, timeouts, and OOMs\n", NumJobs);
 
   GlobalEnv Env;
-  Env.Args = &Args;
+  Env.Args = Args;
+  Env.CorpusDirs = CorpusDirs;
+  Env.Rand = &Rand;
+  Env.Verbosity = Options.Verbosity;
 
   Vector<SizedFile> SeedFiles;
   for (auto &Dir : CorpusDirs)
     GetSizedFilesFromDir(Dir, &SeedFiles);
   std::sort(SeedFiles.begin(), SeedFiles.end());
-  auto TempDir = TempPath(".dir");
-  RmDirRecursive(TempDir);  // just in case there is a leftover from an old run.
-  MkDir(TempDir);
+  Env.TempDir = TempPath(".dir");
+  RmDirRecursive(Env.TempDir);  // in case there is a leftover from old runs.
+  MkDir(Env.TempDir);
 
-  auto CFPath = DirPlusFile(TempDir, "merge.txt");
-  auto LogPath = DirPlusFile(TempDir, "sub.log");
 
   if (CorpusDirs.empty())
-    MkDir(Env.MainCorpusDir = DirPlusFile(TempDir, "C"));
+    MkDir(Env.MainCorpusDir = DirPlusFile(Env.TempDir, "C"));
   else
     Env.MainCorpusDir = CorpusDirs[0];
 
-  auto TempCorpusDir = DirPlusFile(TempDir, "C0");
-
-  CrashResistantMerge(*Env.Args, {}, SeedFiles, &Env.Files, {}, &Env.Features,
+  auto CFPath = DirPlusFile(Env.TempDir, "merge.txt");
+  CrashResistantMerge(Env.Args, {}, SeedFiles, &Env.Files, {}, &Env.Features,
                       CFPath, false);
   RemoveFile(CFPath);
-  Printf("INFO: -fork=1: %zd seeds, starting to fuzz; scratch: %s\n",
-         Env.Files.size(), TempDir.c_str());
-
-  Command BaseCmd(*Env.Args);
-  BaseCmd.removeFlag("fork");
-  for (auto &C : CorpusDirs) // Remove all corpora from the args.
-    BaseCmd.removeArgument(C);
-  BaseCmd.addFlag("reload", "0");  // working in an isolated dir, no reload.
-  int ExitCode = 0;
+  Printf("INFO: -fork=%d: %zd seeds, starting to fuzz; scratch: %s\n",
+         NumJobs, Env.Files.size(), Env.TempDir.c_str());
 
+  int ExitCode = 0;
 
-  for (size_t i = 1; i < 1000000; i++) {
-    // TODO: take new files from disk e.g. those generated by another process.
-
-    FuzzJob Job;
-    Job.Cmd = BaseCmd;
-    if (size_t CorpusSubsetSize = std::min(Env.Files.size(), (size_t)100))
-      for (size_t i = 0; i < CorpusSubsetSize; i++)
-        Job.Files.push_back(Env.Files[Rand.SkewTowardsLast(Env.Files.size())]);
-    Job.CorpusDir = TempCorpusDir;
-    Job.LogPath = LogPath;
-    Job.CFPath = CFPath;
-    // Start from very short runs and gradually increase them.
-    Job.MaxTotalTimeSec = std::min(300, (int)i);
-    RunOneFuzzingJob(&Job);
+  JobQueue FuzzQ, MergeQ;
+  std::atomic<bool> Stop(false);
 
-    if (Options.Verbosity >= 2)
-      Printf("done [%d] %s\n", Job.ExitCode, Job.Cmd.toString().c_str());
-    if (Job.ExitCode == Options.InterruptExitCode)
-      break;
+  size_t JobId = 1;
+  Vector<std::thread> Threads;
+  for (int t = 0; t < NumJobs; t++) {
+    Threads.push_back(std::thread(WorkerThread, &Stop, &FuzzQ, &MergeQ));
+    FuzzQ.Push(Env.CreateNewJob(JobId++));
+  }
 
-    RunOneMergeJob(&Env, &Job);
+  while (!Stop) {
+    auto Job = MergeQ.Pop();
+    if (!Job) {
+      SleepSeconds(1);
+      continue;
+    }
+    ExitCode = Job->ExitCode;
+    if (ExitCode != Options.InterruptExitCode)
+      Env.RunOneMergeJob(Job);
 
     // Continue if our crash is one of the ignorred ones.
-    if (Options.IgnoreTimeouts && Job.ExitCode == Options.TimeoutExitCode)
-      continue;
-    if (Options.IgnoreOOMs && Job.ExitCode == Options.OOMExitCode)
-      continue;
-    // And exit if we don't ignore this crash.
-    if (Job.ExitCode != 0) {
+    if (Options.IgnoreTimeouts && ExitCode == Options.TimeoutExitCode)
+      ;
+    else if (Options.IgnoreOOMs && ExitCode == Options.OOMExitCode)
+      ;
+    else if (ExitCode == Options.InterruptExitCode)
+      Stop = true;
+    else if (ExitCode != 0) {
+      // And exit if we don't ignore this crash.
       Printf("INFO: log from the inner process:\n%s",
-             FileToString(LogPath).c_str());
-      ExitCode = Job.ExitCode;
-      break;
+                   FileToString(Job->LogPath).c_str());
+      Stop = true;
     }
+    RemoveFile(Job->LogPath);
+    delete Job;
+    FuzzQ.Push(Env.CreateNewJob(JobId++));
   }
+  Stop = true;
+
+  for (auto &T : Threads)
+    T.join();
 
-  RmDirRecursive(TempDir);
+  RmDirRecursive(Env.TempDir);
 
   // Use the exit code from the last child process.
   Printf("Fork: exiting: %d\n", ExitCode);
diff --git a/lib/fuzzer/FuzzerFork.h b/lib/fuzzer/FuzzerFork.h
index d9e8de515..b29a43e13 100644
--- a/lib/fuzzer/FuzzerFork.h
+++ b/lib/fuzzer/FuzzerFork.h
@@ -18,7 +18,7 @@
 namespace fuzzer {
 void FuzzWithFork(Random &Rand, const FuzzingOptions &Options,
                   const Vector<std::string> &Args,
-                  const Vector<std::string> &CorpusDirs);
+                  const Vector<std::string> &CorpusDirs, int NumJobs);
 } // namespace fuzzer
 
 #endif // LLVM_FUZZER_FORK_H