Add `--gtest_shard_index` and `--gtest_total_shards` command line arguments.HEADmain

This uses the environment variables as the default value for the command line argument, so this is not a breaking change to most reasonable existing use cases. This *is* a breaking change if a process expects changes to the sharding environment variables to be visible after parsing the command line arguments.

The motivation for this is that some environments do not support environment variables, which prevents the usage of test sharding. These environments also run on simulators and other slow hardware, where sharding is especially important. I've tried a few other things, like hacking together a fake implementation of `getenv`/`setenv`, but this is a pretty unreliable approach on the platform in question.

This also seems like an improvement in consistency to me. Currently, other test selection mechanisms (e.g. filtering, shuffling, repeat), `use arguments (or at least have command line arguments as an option), while as far as I can tell, sharding is the only mechanism which can *only* be specified with environment variables.

PiperOrigin-RevId: 892324928
Change-Id: I5cf814e46e16072e7c160e54c426b02300fe712b

4 files changed, 86 insertions, 79 deletions

diff --git a/googletest/include/gtest/gtest.h b/googletest/include/gtest/gtest.h
index 4218ade..f4a91ee 100644
--- a/googletest/include/gtest/gtest.h
+++ b/googletest/include/gtest/gtest.h
@@ -137,6 +137,10 @@ GTEST_DECLARE_int32_(repeat);
 // only torn down once, for the last.
 GTEST_DECLARE_bool_(recreate_environments_when_repeating);
 
+// Together these flags determine which tests are run if the test is sharded.
+GTEST_DECLARE_int32_(shard_index);
+GTEST_DECLARE_int32_(total_shards);
+
 // This flag controls whether Google Test includes Google Test internal
 // stack frames in failure stack traces.
 GTEST_DECLARE_bool_(show_internal_stack_frames);
diff --git a/googletest/src/gtest-internal-inl.h b/googletest/src/gtest-internal-inl.h
index 6a39b93..2ebfcd3 100644
--- a/googletest/src/gtest-internal-inl.h
+++ b/googletest/src/gtest-internal-inl.h
@@ -246,15 +246,12 @@ GTEST_API_ std::string WideStringToUtf8(const wchar_t* str, int num_chars);
 // be created, prints an error and exits.
 void WriteToShardStatusFileIfNeeded();
 
-// Checks whether sharding is enabled by examining the relevant
-// environment variable values. If the variables are present,
-// but inconsistent (e.g., shard_index >= total_shards), prints
-// an error and exits. If in_subprocess_for_death_test, sharding is
+// Checks whether sharding is enabled by examining the relevant flag values.
+// If the flags are set, but inconsistent (e.g., shard_index >= total_shards),
+// prints an error and exits. If in_subprocess_for_death_test, sharding is
 // disabled because it must only be applied to the original test
 // process. Otherwise, we could filter out death tests we intended to execute.
-GTEST_API_ bool ShouldShard(const char* total_shards_str,
-                            const char* shard_index_str,
-                            bool in_subprocess_for_death_test);
+GTEST_API_ bool ShouldShard(bool in_subprocess_for_death_test);
 
 // Parses the environment variable var as a 32-bit integer. If it is unset,
 // returns default_val. If it is not a 32-bit integer, prints an error and
diff --git a/googletest/src/gtest.cc b/googletest/src/gtest.cc
index 8a38018..ac90786 100644
--- a/googletest/src/gtest.cc
+++ b/googletest/src/gtest.cc
@@ -407,6 +407,18 @@ GTEST_DEFINE_bool_(
     "if exceptions are enabled or exit the program with a non-zero code "
     "otherwise. For use with an external test framework.");
 
+GTEST_DEFINE_int32_(
+    shard_index,
+    testing::internal::Int32FromEnvOrDie(testing::kTestShardIndex, -1),
+    "The zero-based index of the shard to run. A value of -1 "
+    "(the default) indicates that sharding is disabled.");
+
+GTEST_DEFINE_int32_(
+    total_shards,
+    testing::internal::Int32FromEnvOrDie(testing::kTestTotalShards, -1),
+    "The total number of shards to use when running tests in parallel. "
+    "A value of -1 (the default) indicates that sharding is disabled.");
+
 #if GTEST_USE_OWN_FLAGFILE_FLAG_
 GTEST_DEFINE_string_(
     flagfile, testing::internal::StringFromGTestEnv("flagfile", ""),
@@ -3475,11 +3487,11 @@ void PrettyUnitTestResultPrinter::OnTestIterationStart(
                   filter);
   }
 
-  if (internal::ShouldShard(kTestTotalShards, kTestShardIndex, false)) {
-    const int32_t shard_index = Int32FromEnvOrDie(kTestShardIndex, -1);
-    ColoredPrintf(GTestColor::kYellow, "Note: This is test shard %d of %s.\n",
+  if (internal::ShouldShard(false)) {
+    const int32_t shard_index = GTEST_FLAG_GET(shard_index);
+    ColoredPrintf(GTestColor::kYellow, "Note: This is test shard %d of %d.\n",
                   static_cast<int>(shard_index) + 1,
-                  internal::posix::GetEnv(kTestTotalShards));
+                  GTEST_FLAG_GET(total_shards));
   }
 
   if (GTEST_FLAG_GET(shuffle)) {
@@ -5983,8 +5995,7 @@ bool UnitTestImpl::RunAllTests() {
 #endif  // defined(GTEST_EXTRA_DEATH_TEST_CHILD_SETUP_)
 #endif  // GTEST_HAS_DEATH_TEST
 
-  const bool should_shard = ShouldShard(kTestTotalShards, kTestShardIndex,
-                                        in_subprocess_for_death_test);
+  const bool should_shard = ShouldShard(in_subprocess_for_death_test);
 
   // Compares the full test names with the filter to decide which
   // tests to run.
@@ -6196,45 +6207,44 @@ void WriteToShardStatusFileIfNeeded() {
 }
 #endif  // GTEST_HAS_FILE_SYSTEM
 
-// Checks whether sharding is enabled by examining the relevant
-// environment variable values. If the variables are present,
-// but inconsistent (i.e., shard_index >= total_shards), prints
-// an error and exits. If in_subprocess_for_death_test, sharding is
-// disabled because it must only be applied to the original test
-// process. Otherwise, we could filter out death tests we intended to execute.
-bool ShouldShard(const char* total_shards_env, const char* shard_index_env,
-                 bool in_subprocess_for_death_test) {
+// Checks whether sharding is enabled by examining the relevant command line
+// arguments. If the arguments are present, but inconsistent
+// (i.e., shard_index >= total_shards), prints an error and exits.
+// If in_subprocess_for_death_test, sharding is disabled because it must only
+// be applied to the original test process. Otherwise, we could filter out death
+// tests we intended to execute.
+bool ShouldShard(bool in_subprocess_for_death_test) {
   if (in_subprocess_for_death_test) {
     return false;
   }
 
-  const int32_t total_shards = Int32FromEnvOrDie(total_shards_env, -1);
-  const int32_t shard_index = Int32FromEnvOrDie(shard_index_env, -1);
+  const int32_t total_shards = GTEST_FLAG_GET(total_shards);
+  const int32_t shard_index = GTEST_FLAG_GET(shard_index);
 
   if (total_shards == -1 && shard_index == -1) {
     return false;
   } else if (total_shards == -1 && shard_index != -1) {
-    const Message msg = Message() << "Invalid environment variables: you have "
-                                  << kTestShardIndex << " = " << shard_index
-                                  << ", but have left " << kTestTotalShards
-                                  << " unset.\n";
+    const Message msg = Message()
+                        << "Invalid sharding: you have " << kTestShardIndex
+                        << " = " << shard_index << ", but have left "
+                        << kTestTotalShards << " unset.\n";
     ColoredPrintf(GTestColor::kRed, "%s", msg.GetString().c_str());
     fflush(stdout);
     exit(EXIT_FAILURE);
   } else if (total_shards != -1 && shard_index == -1) {
     const Message msg = Message()
-                        << "Invalid environment variables: you have "
-                        << kTestTotalShards << " = " << total_shards
-                        << ", but have left " << kTestShardIndex << " unset.\n";
+                        << "Invalid sharding: you have " << kTestTotalShards
+                        << " = " << total_shards << ", but have left "
+                        << kTestShardIndex << " unset.\n";
     ColoredPrintf(GTestColor::kRed, "%s", msg.GetString().c_str());
     fflush(stdout);
     exit(EXIT_FAILURE);
   } else if (shard_index < 0 || shard_index >= total_shards) {
     const Message msg =
-        Message() << "Invalid environment variables: we require 0 <= "
-                  << kTestShardIndex << " < " << kTestTotalShards
-                  << ", but you have " << kTestShardIndex << "=" << shard_index
-                  << ", " << kTestTotalShards << "=" << total_shards << ".\n";
+        Message() << "Invalid sharding: we require 0 <= " << kTestShardIndex
+                  << " < " << kTestTotalShards << ", but you have "
+                  << kTestShardIndex << "=" << shard_index << ", "
+                  << kTestTotalShards << "=" << total_shards << ".\n";
     ColoredPrintf(GTestColor::kRed, "%s", msg.GetString().c_str());
     fflush(stdout);
     exit(EXIT_FAILURE);
@@ -6277,11 +6287,10 @@ bool ShouldRunTestOnShard(int total_shards, int shard_index, int test_id) {
 // . Returns the number of tests that should run.
 int UnitTestImpl::FilterTests(ReactionToSharding shard_tests) {
   const int32_t total_shards = shard_tests == HONOR_SHARDING_PROTOCOL
-                                   ? Int32FromEnvOrDie(kTestTotalShards, -1)
+                                   ? GTEST_FLAG_GET(total_shards)
                                    : -1;
-  const int32_t shard_index = shard_tests == HONOR_SHARDING_PROTOCOL
-                                  ? Int32FromEnvOrDie(kTestShardIndex, -1)
-                                  : -1;
+  const int32_t shard_index =
+      shard_tests == HONOR_SHARDING_PROTOCOL ? GTEST_FLAG_GET(shard_index) : -1;
 
   const PositiveAndNegativeUnitTestFilter gtest_flag_filter(
       GTEST_FLAG_GET(filter));
@@ -6810,6 +6819,8 @@ static bool ParseGoogleTestFlag(const char* const arg) {
   GTEST_INTERNAL_PARSE_FLAG(print_utf8);
   GTEST_INTERNAL_PARSE_FLAG(random_seed);
   GTEST_INTERNAL_PARSE_FLAG(repeat);
+  GTEST_INTERNAL_PARSE_FLAG(shard_index);
+  GTEST_INTERNAL_PARSE_FLAG(total_shards);
   GTEST_INTERNAL_PARSE_FLAG(recreate_environments_when_repeating);
   GTEST_INTERNAL_PARSE_FLAG(shuffle);
   GTEST_INTERNAL_PARSE_FLAG(stack_trace_depth);
diff --git a/googletest/test/gtest_unittest.cc b/googletest/test/gtest_unittest.cc
index 5304c3a..8ce64d3 100644
--- a/googletest/test/gtest_unittest.cc
+++ b/googletest/test/gtest_unittest.cc
@@ -1873,36 +1873,31 @@ TEST(ShouldRunTestOnShardTest, IsPartitionWhenThereIsOneShard) {
 
 class ShouldShardTest : public testing::Test {
  protected:
-  void SetUp() override {
-    index_var_ = GTEST_FLAG_PREFIX_UPPER_ "INDEX";
-    total_var_ = GTEST_FLAG_PREFIX_UPPER_ "TOTAL";
-  }
+  void SetUp() override {}
 
   void TearDown() override {
-    SetEnv(index_var_, "");
-    SetEnv(total_var_, "");
+    GTEST_FLAG_SET(shard_index, -1);
+    GTEST_FLAG_SET(total_shards, -1);
   }
-
-  const char* index_var_;
-  const char* total_var_;
 };
 
 // Tests that sharding is disabled if neither of the environment variables
 // are set.
 TEST_F(ShouldShardTest, ReturnsFalseWhenNeitherEnvVarIsSet) {
-  SetEnv(index_var_, "");
-  SetEnv(total_var_, "");
+  GTEST_FLAG_SET(shard_index, -1);
+  GTEST_FLAG_SET(total_shards, -1);
 
-  EXPECT_FALSE(ShouldShard(total_var_, index_var_, false));
-  EXPECT_FALSE(ShouldShard(total_var_, index_var_, true));
+  EXPECT_FALSE(ShouldShard(false));
+  EXPECT_FALSE(ShouldShard(true));
 }
 
 // Tests that sharding is not enabled if total_shards  == 1.
 TEST_F(ShouldShardTest, ReturnsFalseWhenTotalShardIsOne) {
-  SetEnv(index_var_, "0");
-  SetEnv(total_var_, "1");
-  EXPECT_FALSE(ShouldShard(total_var_, index_var_, false));
-  EXPECT_FALSE(ShouldShard(total_var_, index_var_, true));
+  GTEST_FLAG_SET(shard_index, 0);
+  GTEST_FLAG_SET(total_shards, 1);
+
+  EXPECT_FALSE(ShouldShard(false));
+  EXPECT_FALSE(ShouldShard(true));
 }
 
 // Tests that sharding is enabled if total_shards > 1 and
@@ -1910,20 +1905,20 @@ TEST_F(ShouldShardTest, ReturnsFalseWhenTotalShardIsOne) {
 // Environment variables are not supported on Windows CE.
 #ifndef GTEST_OS_WINDOWS_MOBILE
 TEST_F(ShouldShardTest, WorksWhenShardEnvVarsAreValid) {
-  SetEnv(index_var_, "4");
-  SetEnv(total_var_, "22");
-  EXPECT_TRUE(ShouldShard(total_var_, index_var_, false));
-  EXPECT_FALSE(ShouldShard(total_var_, index_var_, true));
+  GTEST_FLAG_SET(shard_index, 4);
+  GTEST_FLAG_SET(total_shards, 22);
+  EXPECT_TRUE(ShouldShard(false));
+  EXPECT_FALSE(ShouldShard(true));
 
-  SetEnv(index_var_, "8");
-  SetEnv(total_var_, "9");
-  EXPECT_TRUE(ShouldShard(total_var_, index_var_, false));
-  EXPECT_FALSE(ShouldShard(total_var_, index_var_, true));
+  GTEST_FLAG_SET(shard_index, 8);
+  GTEST_FLAG_SET(total_shards, 9);
+  EXPECT_TRUE(ShouldShard(false));
+  EXPECT_FALSE(ShouldShard(true));
 
-  SetEnv(index_var_, "0");
-  SetEnv(total_var_, "9");
-  EXPECT_TRUE(ShouldShard(total_var_, index_var_, false));
-  EXPECT_FALSE(ShouldShard(total_var_, index_var_, true));
+  GTEST_FLAG_SET(shard_index, 0);
+  GTEST_FLAG_SET(total_shards, 9);
+  EXPECT_TRUE(ShouldShard(false));
+  EXPECT_FALSE(ShouldShard(true));
 }
 #endif  // !GTEST_OS_WINDOWS_MOBILE
 
@@ -1932,21 +1927,21 @@ TEST_F(ShouldShardTest, WorksWhenShardEnvVarsAreValid) {
 typedef ShouldShardTest ShouldShardDeathTest;
 
 TEST_F(ShouldShardDeathTest, AbortsWhenShardingEnvVarsAreInvalid) {
-  SetEnv(index_var_, "4");
-  SetEnv(total_var_, "4");
-  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(total_var_, index_var_, false), ".*");
+  GTEST_FLAG_SET(shard_index, 4);
+  GTEST_FLAG_SET(total_shards, 4);
+  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(false), ".*");
 
-  SetEnv(index_var_, "4");
-  SetEnv(total_var_, "-2");
-  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(total_var_, index_var_, false), ".*");
+  GTEST_FLAG_SET(shard_index, 4);
+  GTEST_FLAG_SET(total_shards, -2);
+  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(false), ".*");
 
-  SetEnv(index_var_, "5");
-  SetEnv(total_var_, "");
-  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(total_var_, index_var_, false), ".*");
+  GTEST_FLAG_SET(shard_index, 5);
+  GTEST_FLAG_SET(total_shards, 5);
+  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(false), ".*");
 
-  SetEnv(index_var_, "");
-  SetEnv(total_var_, "5");
-  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(total_var_, index_var_, false), ".*");
+  GTEST_FLAG_SET(shard_index, -1);
+  GTEST_FLAG_SET(total_shards, 5);
+  EXPECT_DEATH_IF_SUPPORTED(ShouldShard(false), ".*");
 }
 
 // Tests that ShouldRunTestOnShard is a partition when 5