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

…Centipede. PiperOrigin-RevId: 689060275
google · Oct 23, 2024 · 551bb67 · 551bb67
1 parent 19fd030
commit 551bb67
Show file tree

Hide file tree

Showing 9 changed files with 52 additions and 5 deletions.
diff --git a/centipede/BUILD b/centipede/BUILD
@@ -741,6 +741,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
         "@com_google_fuzztest//common:logging",
+        "@com_google_fuzztest//fuzztest:configuration",
     ],
 )
 
@@ -812,6 +813,7 @@ cc_library(
         ":distill",
         ":early_exit",
         ":environment",
+        ":environment_flags",
         ":minimize_crash",
         ":pc_info",
         ":periodic_action",

diff --git a/centipede/centipede_interface.cc b/centipede/centipede_interface.cc
@@ -53,6 +53,7 @@
 #include "./centipede/distill.h"
 #include "./centipede/early_exit.h"
 #include "./centipede/environment.h"
+#include "./centipede/environment_flags.h"
 #include "./centipede/minimize_crash.h"
 #include "./centipede/pc_info.h"
 #include "./centipede/periodic_action.h"
@@ -674,8 +675,9 @@ int CentipedeMain(const Environment &env,
                "fuzz test.";
         CHECK(time_limit_per_test >= absl::Seconds(1))
             << "Time limit per fuzz test must be at least 1 second.";
-        return UpdateCorpusDatabaseForFuzzTests(env, *target_config,
-                                                callbacks_factory);
+        return UpdateCorpusDatabaseForFuzzTests(
+            AdjustEnvironmentForTargetConfig(env, *target_config),
+            *target_config, callbacks_factory);
       }
     } else if (std::getenv("CENTIPEDE_NO_FUZZ_IF_NO_CONFIG") != nullptr) {
       // Target config is empty when the shard does not contain any fuzz tests.

diff --git a/centipede/environment_flags.cc b/centipede/environment_flags.cc
@@ -32,6 +32,7 @@
 #include "./centipede/environment.h"
 #include "./centipede/util.h"
 #include "./common/logging.h"
+#include "./fuzztest/internal/configuration.h"
 
 namespace {
 const auto *default_env = new centipede::Environment();
@@ -572,4 +573,15 @@ Environment CreateEnvironmentFromFlags(const std::vector<std::string> &argv) {
   return env_from_flags;
 }
 
+Environment AdjustEnvironmentForTargetConfig(
+    Environment env, const fuzztest::internal::Configuration &config) {
+  if (config.jobs != 0) {
+    CHECK(absl::GetFlag(FLAGS_j) == 0 || absl::GetFlag(FLAGS_j) == config.jobs);
+    env.total_shards = config.jobs;
+    env.num_threads = config.jobs;
+    env.my_shard_index = 0;
+  }
+  return env;
+}
+
 }  // namespace centipede
diff --git a/centipede/environment_flags.h b/centipede/environment_flags.h
@@ -19,6 +19,7 @@
 #include <vector>
 
 #include "./centipede/environment.h"
+#include "./fuzztest/internal/configuration.h"
 
 namespace centipede {
 
@@ -27,6 +28,12 @@ namespace centipede {
 Environment CreateEnvironmentFromFlags(
     const std::vector<std::string> &argv = {});
 
+// Returns `env` adjusted for the `config` obtained from the target binary.
+// Check-fails if the values in `config` are inconsistent with the corresponding
+// values passed by flags.
+Environment AdjustEnvironmentForTargetConfig(
+    Environment env, const fuzztest::internal::Configuration &config);
+
 }  // namespace centipede
 
 #endif  // THIRD_PARTY_CENTIPEDE_ENVIRONMENT_FLAGS_H_
diff --git a/e2e_tests/corpus_database_test.cc b/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("j", "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
diff --git a/fuzztest/init_fuzztest.cc b/fuzztest/init_fuzztest.cc
@@ -160,6 +160,11 @@ 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(size_t, j, 0,
+                     "When running with Centipede, the number of fuzzing jobs "
+                     "to run in parallel. If 0, Centipede determines the "
+                     "number of jobs from its own flags.");
+
 namespace fuzztest {
 
 std::vector<std::string> ListRegisteredTests() {
@@ -255,7 +260,8 @@ 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=*/absl::GetFlag(FUZZTEST_FLAG(j))};
 }
 }  // namespace
 

diff --git a/fuzztest/internal/configuration.cc b/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)};
   }();

diff --git a/fuzztest/internal/configuration.h b/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;

diff --git a/fuzztest/internal/configuration_test.cc b/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"};