Add a FuzzTest flag for the number of fuzzing jobs when running with …

…Centipede.

PiperOrigin-RevId: 690797852

centipede/BUILD

@@ -856,6 +856,7 @@ cc_library(
     deps = [
         ":feature",
         ":knobs",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
@@ -865,6 +866,7 @@ cc_library(
         "@com_google_fuzztest//common:logging",
         "@com_google_fuzztest//common:remote_file",
         "@com_google_fuzztest//common:status_macros",
+        "@com_google_fuzztest//fuzztest:configuration",
     ],
 )
 

centipede/centipede_interface.cc

@@ -474,6 +474,8 @@ SeedCorpusConfig GetSeedCorpusConfig(const Environment &env,
 int UpdateCorpusDatabaseForFuzzTests(
     Environment env, const fuzztest::internal::Configuration &fuzztest_config,
     CentipedeCallbacksFactory &callbacks_factory) {
+  env.UpdateWithTargetConfig(fuzztest_config);
+
   absl::Time start_time = absl::Now();
   LOG(INFO) << "Starting the update of the corpus database for fuzz tests:"
             << "\nBinary: " << env.binary

centipede/environment.cc

@@ -23,6 +23,7 @@
 #include <system_error>  // NOLINT
 #include <vector>
 
+#include "absl/base/no_destructor.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
@@ -34,9 +35,15 @@
 #include "./common/logging.h"
 #include "./common/remote_file.h"
 #include "./common/status_macros.h"
+#include "./fuzztest/internal/configuration.h"
 
 namespace centipede {
 
+const Environment &Environment::Default() {
+  static absl::NoDestructor<Environment> default_env;
+  return *default_env;
+}
+
 bool Environment::DumpCorpusTelemetryInThisShard() const {
   // Corpus stats are global across all shards on all machines.
   return my_shard_index == 0 && telemetry_frequency != 0;
@@ -205,4 +212,17 @@ void Environment::ReadKnobsFileIfSpecified() {
   });
 }
 
+void Environment::UpdateWithTargetConfig(
+    const fuzztest::internal::Configuration &config) {
+  if (config.jobs == 0) return;
+  CHECK(j == Default().j || j == config.jobs)
+      << "Value for --j is inconsistent with the value for jobs in the target "
+         "binary:"
+      << VV(j) << VV(config.jobs);
+  j = config.jobs;
+  total_shards = config.jobs;
+  num_threads = config.jobs;
+  my_shard_index = 0;
+}
+
 }  // namespace centipede

centipede/environment.h

@@ -26,6 +26,7 @@
 #include "absl/time/time.h"
 #include "./centipede/feature.h"
 #include "./centipede/knobs.h"
+#include "./fuzztest/internal/configuration.h"
 
 namespace centipede {
 
@@ -46,6 +47,7 @@ struct Environment {
   size_t total_shards = 1;
   size_t my_shard_index = 0;
   size_t num_threads = 1;
+  size_t j = 0;
   size_t max_len = 4000;
   size_t batch_size = 1000;
   size_t mutate_batch_size = 2;
@@ -152,6 +154,9 @@ struct Environment {
 
   // APIs ----------------------------------------------------------------------
 
+  // Returns an instance of the environment with default values.
+  static const Environment& Default();
+
   // Should certain actions be performed ---------------------------------------
 
   // Returns true if we want to log features as symbols in this shard.
@@ -205,6 +210,10 @@ struct Environment {
 
   // Reads `knobs` from `knobs_file`. Does nothing if the `knobs_file` is empty.
   void ReadKnobsFileIfSpecified();
+  // Updates `this` with `config` obtained from the target binary. CHECK-fails
+  // if the fields are non-default and inconsistent with the corresponding
+  // values in `config`.
+  void UpdateWithTargetConfig(const fuzztest::internal::Configuration& config);
 };
 
 }  // namespace centipede

e2e_tests/corpus_database_test.cc

@@ -55,7 +55,8 @@ class UpdateCorpusDatabaseTest : public testing::Test {
                            " ",
                            CreateFuzzTestFlag("corpus_database",
                                               GetCorpusDatabasePath()),
-                           " ", CreateFuzzTestFlag("fuzz_for", "30s"))}}});
+                           " ", CreateFuzzTestFlag("fuzz_for", "30s"), " ",
+                           CreateFuzzTestFlag("jobs", "2"))}}});
 
     *centipede_std_out_ = std::move(std_out);
     *centipede_std_err_ = std::move(std_err);
@@ -90,5 +91,13 @@ TEST_F(UpdateCorpusDatabaseTest, RunsFuzzTests) {
               HasSubstr("Fuzzing FuzzTest.FailsInTwoWays"));
 }
 
+TEST_F(UpdateCorpusDatabaseTest, UsesMultipleShardsForFuzzingAndDistillation) {
+  EXPECT_THAT(
+      GetCentipedeStdErr(),
+      AllOf(HasSubstr("[S0.0] begin-fuzz"), HasSubstr("[S1.0] begin-fuzz"),
+            HasSubstr("DISTILL[S.0]: Distilling to output shard 0"),
+            HasSubstr("DISTILL[S.1]: Distilling to output shard 1")));
+}
+
 }  // namespace
 }  // namespace fuzztest::internal

fuzztest/init_fuzztest.cc

@@ -159,6 +159,10 @@ FUZZTEST_DEFINE_FLAG(
     "for an input if the execution of the property-function with the input "
     "takes longer than this time limit.");
 
+FUZZTEST_DEFINE_FLAG(std::optional<size_t>, jobs, std::nullopt,
+                     "The number of fuzzing jobs to run in parallel. If "
+                     "unspecified, the number of jobs is 1.");
+
 namespace fuzztest {
 
 std::vector<std::string> ListRegisteredTests() {
@@ -244,6 +248,9 @@ internal::Configuration CreateConfigurationsFromFlags(
                               : replay_corpus_time_limit
                                   ? *replay_corpus_time_limit
                                   : absl::ZeroDuration();
+  std::optional<size_t> jobs = absl::GetFlag(FUZZTEST_FLAG(jobs));
+  FUZZTEST_INTERNAL_CHECK(!jobs.has_value() || *jobs > 0, "If specified, --",
+                          FUZZTEST_FLAG(jobs).Name(), " must be positive.");
   return internal::Configuration{
       absl::GetFlag(FUZZTEST_FLAG(corpus_database)),
       /*stats_root=*/"",
@@ -254,7 +261,7 @@ internal::Configuration CreateConfigurationsFromFlags(
       /*stack_limit=*/absl::GetFlag(FUZZTEST_FLAG(stack_limit_kb)) * 1024,
       /*rss_limit=*/absl::GetFlag(FUZZTEST_FLAG(rss_limit_mb)) * 1024 * 1024,
       absl::GetFlag(FUZZTEST_FLAG(time_limit_per_input)), time_limit,
-      absl::GetFlag(FUZZTEST_FLAG(time_budget_type))};
+      absl::GetFlag(FUZZTEST_FLAG(time_budget_type)), jobs.value_or(0)};
 }
 }  // namespace
 

fuzztest/internal/configuration.cc

@@ -208,7 +208,7 @@ std::string Configuration::Serialize() const {
              SpaceFor(reproduce_findings_as_separate_tests) +
              SpaceFor(stack_limit) + SpaceFor(rss_limit) +
              SpaceFor(time_limit_per_input_str) + SpaceFor(time_limit_str) +
-             SpaceFor(time_budget_type_str) +
+             SpaceFor(time_budget_type_str) + SpaceFor(jobs) +
              SpaceFor(crashing_input_to_reproduce) +
              SpaceFor(reproduction_command_template));
   size_t offset = 0;
@@ -223,6 +223,7 @@ std::string Configuration::Serialize() const {
   offset = WriteString(out, offset, time_limit_per_input_str);
   offset = WriteString(out, offset, time_limit_str);
   offset = WriteString(out, offset, time_budget_type_str);
+  offset = WriteIntegral(out, offset, jobs);
   offset = WriteOptionalString(out, offset, crashing_input_to_reproduce);
   offset = WriteOptionalString(out, offset, reproduction_command_template);
   CHECK_EQ(offset, out.size());
@@ -245,6 +246,7 @@ absl::StatusOr<Configuration> Configuration::Deserialize(
     ASSIGN_OR_RETURN(time_limit_per_input_str, ConsumeString(serialized));
     ASSIGN_OR_RETURN(time_limit_str, ConsumeString(serialized));
     ASSIGN_OR_RETURN(time_budget_type_str, ConsumeString(serialized));
+    ASSIGN_OR_RETURN(jobs, Consume<size_t>(serialized));
     ASSIGN_OR_RETURN(crashing_input_to_reproduce,
                      ConsumeOptionalString(serialized));
     ASSIGN_OR_RETURN(reproduction_command_template,
@@ -269,6 +271,7 @@ absl::StatusOr<Configuration> Configuration::Deserialize(
                          *time_limit_per_input,
                          *time_limit,
                          *time_budget_type,
+                         *jobs,
                          *std::move(crashing_input_to_reproduce),
                          *std::move(reproduction_command_template)};
   }();

fuzztest/internal/configuration.h

@@ -70,6 +70,9 @@ struct Configuration {
   absl::Duration time_limit = absl::InfiniteDuration();
   // Whether the time limit is for each test or for all tests in the binary.
   TimeBudgetType time_budget_type = TimeBudgetType::kPerTest;
+  // The number of fuzzing jobs to run in parallel. Zero indicates that the
+  // number of jobs is unspecified by the test binary.
+  size_t jobs = 0;
 
   // When set, `FuzzTestFuzzer` replays only one input (no fuzzing is done).
   std::optional<std::string> crashing_input_to_reproduce;

fuzztest/internal/configuration_test.cc

@@ -26,6 +26,7 @@ MATCHER_P(IsOkAndEquals, config, "") {
          config.time_limit_per_input == other->time_limit_per_input &&
          config.time_limit == other->time_limit &&
          config.time_budget_type == other->time_budget_type &&
+         config.jobs == other->jobs &&
          config.crashing_input_to_reproduce ==
              other->crashing_input_to_reproduce &&
          config.reproduction_command_template ==
@@ -45,6 +46,7 @@ TEST(ConfigurationTest,
                               /*time_limit_per_input=*/absl::Seconds(42),
                               /*time_limit=*/absl::Minutes(42),
                               /*time_budget_type=*/TimeBudgetType::kPerTest,
+                              /*jobs=*/1,
                               /*crashing_input_to_reproduce=*/std::nullopt,
                               /*reproduction_command_template=*/std::nullopt};
 
@@ -65,6 +67,7 @@ TEST(ConfigurationTest,
                               /*time_limit_per_input=*/absl::Seconds(42),
                               /*time_limit=*/absl::Minutes(42),
                               /*time_budget_type=*/TimeBudgetType::kPerTest,
+                              /*jobs=*/1,
                               "crashing_input_to_reproduce",
                               "reproduction_command_template"};