Add support for having different default benchmarks in single-threaded and multi-threaded runs.

PiperOrigin-RevId: 806013441
Change-Id: If823a13a45931122ee1a0d01e35f540cfd171887

Author: liyuying0000

fleetbench/BUILD

@@ -117,7 +117,6 @@ cc_library(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/log:globals",
         "@com_google_absl//absl/log:initialize",
-        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_benchmark//:benchmark",

fleetbench/benchmark_main.cc

@@ -129,11 +129,17 @@ int main(int argc, char* argv[]) {
 
   CheckMemoryLimit();
 
-  if (benchmark::GetBenchmarkFilter().empty() ||
-      benchmark::GetBenchmarkFilter() == "default") {
-    // --benchmark_filter flag not set
+  const std::string benchmark_filter = benchmark::GetBenchmarkFilter();
+  if (benchmark_filter.empty() || benchmark_filter == "default" ||
+      benchmark_filter == "single") {
+    // --benchmark_filter flag not set, set to default or "single"
     benchmark::SetBenchmarkFilter(
-        fleetbench::DynamicRegistrar::Get()->GetDefaultFilter());
+        fleetbench::DynamicRegistrar::Get()->GetDefaultFilter(
+            fleetbench::ThreadingMode::kSingleThreaded));
+  } else if (benchmark_filter == "multi") {
+    benchmark::SetBenchmarkFilter(
+        fleetbench::DynamicRegistrar::Get()->GetDefaultFilter(
+            fleetbench::ThreadingMode::kMultiThreaded));
   }
   fleetbench::DynamicRegistrar::Get()->Run();
 

fleetbench/dynamic_registrar.cc

@@ -36,11 +36,27 @@ void DynamicRegistrar::Run() {
   callbacks_.clear();
 }
 
-void DynamicRegistrar::AddDefaultFilter(std::string filter) {
-  if (!default_filter_.empty()) {
-    default_filter_.append("|");
+void DynamicRegistrar::AddDefaultFilter(std::string filter,
+                                        ThreadingMode mode) {
+  if (mode == ThreadingMode::kSingleThreaded || mode == ThreadingMode::kBoth) {
+    if (!default_single_threaded_filter_.empty()) {
+      default_single_threaded_filter_.push_back('|');
+    }
+    default_single_threaded_filter_.append(filter);
   }
-  default_filter_.append(filter);
+  if (mode == ThreadingMode::kMultiThreaded || mode == ThreadingMode::kBoth) {
+    if (!default_multi_threaded_filter_.empty()) {
+      default_multi_threaded_filter_.push_back('|');
+    }
+    default_multi_threaded_filter_.append(filter);
+  }
+}
+
+std::string DynamicRegistrar::GetDefaultFilter(ThreadingMode mode) const {
+  if (mode == ThreadingMode::kSingleThreaded) {
+    return default_single_threaded_filter_;
+  }
+  return default_multi_threaded_filter_;
 }
 
 }  // namespace fleetbench

fleetbench/dynamic_registrar.h

@@ -21,6 +21,12 @@
 
 namespace fleetbench {
 
+enum class ThreadingMode {
+  kSingleThreaded,
+  kMultiThreaded,
+  kBoth,
+};
+
 #define REGISTER_BENCHMARK(n) benchmark::RegisterBenchmark(#n, n)
 #define REGISTER_BENCHMARK_TEMPLATE(n, ...) \
   benchmark::RegisterBenchmark(#n "<" #__VA_ARGS__ ">", n<__VA_ARGS__>)
@@ -45,15 +51,17 @@ class DynamicRegistrar {
 
   // Used by fleetbench main() as the default benchmark filter if the
   // --benchmark_filter flag is not set. The returned filter is the disjunction
-  // of all filters that were added by calls to AddDefaultFilter.
-  std::string GetDefaultFilter() const { return default_filter_; }
+  // of all filters that were added by calls to AddDefaultFilter for the given
+  // `ThreadingMode`.
+  std::string GetDefaultFilter(ThreadingMode mode) const;
 
-  // Adds an additional benchmark filter.
-  void AddDefaultFilter(std::string filter);
+  void AddDefaultFilter(std::string filter,
+                        ThreadingMode mode = ThreadingMode::kBoth);
 
  private:
   std::vector<std::function<void()>> callbacks_;
-  std::string default_filter_;
+  std::string default_single_threaded_filter_;
+  std::string default_multi_threaded_filter_;
 };
 }  // namespace fleetbench
 

fleetbench/dynamic_registrar_test.cc

@@ -33,5 +33,28 @@ TEST(DynamicRegistrar, CallbackRunsOnlyOnce) {
   DynamicRegistrar::Get()->Run();
 }
 
+TEST(DynamicRegistrar, DefaultFilters) {
+  DynamicRegistrar registrar;
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kSingleThreaded), "");
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kMultiThreaded), "");
+
+  registrar.AddDefaultFilter("filter1", ThreadingMode::kSingleThreaded);
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kSingleThreaded),
+            "filter1");
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kMultiThreaded), "");
+
+  registrar.AddDefaultFilter("filter2", ThreadingMode::kMultiThreaded);
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kSingleThreaded),
+            "filter1");
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kMultiThreaded),
+            "filter2");
+
+  registrar.AddDefaultFilter("filter3", ThreadingMode::kBoth);
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kSingleThreaded),
+            "filter1|filter3");
+  EXPECT_EQ(registrar.GetDefaultFilter(ThreadingMode::kMultiThreaded),
+            "filter2|filter3");
+}
+
 }  // namespace
 }  // namespace fleetbench