diff --git a/tcmalloc/BUILD b/tcmalloc/BUILD index db9fd9b3f..1217d6e7e 100644 --- a/tcmalloc/BUILD +++ b/tcmalloc/BUILD @@ -57,6 +57,7 @@ cc_library( "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/crc:crc32c", "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/hash", "@com_google_absl//absl/strings", diff --git a/tcmalloc/experiment.cc b/tcmalloc/experiment.cc index 7739b9562..0acbb3442 100644 --- a/tcmalloc/experiment.cc +++ b/tcmalloc/experiment.cc @@ -14,7 +14,7 @@ #include "tcmalloc/experiment.h" -#include +#include #include #include @@ -27,10 +27,10 @@ #include "absl/base/attributes.h" #include "absl/base/call_once.h" #include "absl/base/internal/cycleclock.h" +#include "absl/crc/crc32c.h" #include "absl/functional/function_ref.h" #include "absl/hash/hash.h" #include "absl/strings/match.h" -#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "tcmalloc/experiment_config.h" #include "tcmalloc/internal/config.h" @@ -75,11 +75,18 @@ const bool* GetSelectedExperiments() { const char* test_target = thread_safe_getenv("TEST_TARGET"); const char* active_experiments = thread_safe_getenv(kExperiments); const char* disabled_experiments = thread_safe_getenv(kDisableExperiments); + struct utsname uts; + absl::string_view hostname; + if (const char* env = thread_safe_getenv("HOSTNAME")) { + hostname = env; + } + SelectExperiments( by_id, test_target ? test_target : "", active_experiments ? active_experiments : "", disabled_experiments ? disabled_experiments : "", - active_experiments == nullptr && disabled_experiments == nullptr); + active_experiments == nullptr && disabled_experiments == nullptr, + hostname); }); return by_id; } @@ -104,16 +111,46 @@ void ParseExperiments(absl::string_view labels, F f) { } // namespace +uint16_t CalculateRolloutBucket(absl::string_view hostname, + absl::string_view salt) { + uint32_t hash = static_cast( + absl::ExtendCrc32c(absl::ComputeCrc32c(hostname), salt)); + hash = hash ^ (hash >> 16); + return hash % 1024; +} + +bool IsExperimentRolloutEnabled(const ExperimentConfig& config, + absl::string_view hostname) { + if (hostname.empty()) { + return false; + } + + // Ensure experiments are i.i.d. from one another by using their own names as + // a salt, unless explicitly requested. + absl::string_view salt = + config.rollout_salt.empty() ? config.name : config.rollout_salt; + uint16_t val = CalculateRolloutBucket(hostname, actual_salt); + return val >= config.rollout_lower_bound && val < config.rollout_upper_bound; +} + const bool* SelectExperiments(bool* buffer, absl::string_view test_target, absl::string_view active, absl::string_view disabled, - bool unset) { + bool unset, absl::string_view hostname) { memset(buffer, 0, sizeof(*buffer) * kNumExperiments); if (active == kEnableAll) { std::fill(buffer, buffer + kNumExperiments, true); } + for (const auto& config : experiments) { + if (config.rollout_upper_bound > 0) { + if (IsExperimentRolloutEnabled(config, hostname)) { + buffer[static_cast(config.id)] = true; + } + } + } + ParseExperiments(active, [buffer](absl::string_view token) { Experiment id; if (LookupExperimentID(token, &id)) { diff --git a/tcmalloc/experiment.h b/tcmalloc/experiment.h index 222dda61d..5601e3544 100644 --- a/tcmalloc/experiment.h +++ b/tcmalloc/experiment.h @@ -54,7 +54,13 @@ constexpr size_t kNumExperiments = const bool* absl_nonnull SelectExperiments( bool* absl_nonnull buffer, absl::string_view test_target, absl::string_view active, absl::string_view disabled, - bool unset); + bool unset, absl::string_view hostname); + +uint16_t CalculateRolloutBucket(absl::string_view hostname, + absl::string_view salt); + +bool IsExperimentRolloutEnabled(const ExperimentConfig& config, + absl::string_view hostname); } // namespace tcmalloc_internal diff --git a/tcmalloc/experiment_config.h b/tcmalloc/experiment_config.h index 6af66b2d2..332a48a79 100644 --- a/tcmalloc/experiment_config.h +++ b/tcmalloc/experiment_config.h @@ -15,6 +15,8 @@ #ifndef TCMALLOC_EXPERIMENT_CONFIG_H_ #define TCMALLOC_EXPERIMENT_CONFIG_H_ +#include + #include "absl/strings/string_view.h" // Autogenerated by experiments_proto_test --experiments_generate_config=true @@ -45,6 +47,9 @@ struct ExperimentConfig { absl::string_view name; bool brittle = false; bool force_disable = false; + uint16_t rollout_lower_bound = 0; + uint16_t rollout_upper_bound = 0; + absl::string_view rollout_salt; }; // clang-format off diff --git a/tcmalloc/experiment_fuzz.cc b/tcmalloc/experiment_fuzz.cc index 53534568c..d3106705d 100644 --- a/tcmalloc/experiment_fuzz.cc +++ b/tcmalloc/experiment_fuzz.cc @@ -19,13 +19,14 @@ #include "fuzztest/fuzztest.h" #include "absl/strings/string_view.h" #include "tcmalloc/experiment.h" +#include "tcmalloc/experiment_config.h" namespace tcmalloc::tcmalloc_internal { namespace { void FuzzSelectExperiments(absl::string_view test_target, absl::string_view active, absl::string_view disabled, - bool unset) { + bool unset, absl::string_view hostname) { if (unset && !test_target.empty() && (!active.empty() || !disabled.empty())) { return; } @@ -33,17 +34,24 @@ void FuzzSelectExperiments(absl::string_view test_target, bool buffer[tcmalloc::tcmalloc_internal::kNumExperiments]; SelectExperiments(buffer, test_target, active, disabled, - unset); + unset, hostname); } FUZZ_TEST(ExperimentTest, FuzzSelectExperiments); +void FuzzRolloutEnabled(const ExperimentConfig& config, + absl::string_view hostname) { + IsExperimentRolloutEnabled(config, hostname); +} + +FUZZ_TEST(ExperimentTest, FuzzRolloutEnabled); + TEST(ExperimentTest, FuzzSelectExperiments_b395212979) { FuzzSelectExperiments( "t_fuenchmark&" "vvvvvvvvvvvvvvvvvvVvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv", "", "", - true); + true, "some_hostname"); } } // namespace