Fix lifetime issue in CompilationEnvironmentsTest

`CompilationEnvironments::RegisterProcessNewEnvFn` requires the passed in `Descriptor*` pointer to be valid for the lifetime of the program.

The test `GetEnvTriggersFullNameFallback` dynamically registers a processing function for a custom descriptor and achieves the lifetime requirements by leaking the `DescriptorPool` instance (call to operator new).

This change fixes that properly by adding a deregister function to `CompilationEnvironments` and deregisters the processing function at the end of the test case.

It also removes all usages of `tsl::protobuf` and replaces them by `proto2`.

PiperOrigin-RevId: 853274937

Author: beckerhe

xla/service/BUILD

@@ -5530,14 +5530,13 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_protobuf//:any_cc_proto",
+        "@com_google_protobuf//:protobuf",
         "@tsl//tsl/platform:logging",  # fixdeps: keep
-        "@tsl//tsl/platform:protobuf",
     ],
 )
 
@@ -5566,16 +5565,18 @@ xla_cc_test(
     deps = [
         ":compilation_environments",
         ":test_compilation_environment_proto_cc",
+        # "@com_google_protobuf//:descriptor_cc_proto" - This target is included in @com_google_protobuf//:protobuf and not available separately in OSS.
+        "@com_google_googletest//:gtest",
+        "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_protobuf//:protobuf",
         "//xla:xla_proto_cc",
         "//xla/hlo/testlib:test",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_matchers",
-        "@com_google_googletest//:gtest",
         "@tsl//tsl/platform:casts",
-        "@tsl//tsl/platform:protobuf",
     ],
 )
 

xla/service/compilation_environments.cc

@@ -28,22 +28,23 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
-#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
+#include "google/protobuf/descriptor.h"
+#include "google/protobuf/message.h"
+#include "google/protobuf/unknown_field_set.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
-#include "tsl/platform/protobuf.h"
 
 namespace xla {
 namespace {
 
 ABSL_CONST_INIT absl::Mutex process_new_env_fns_mu(absl::kConstInit);
-absl::flat_hash_map<const tsl::protobuf::Descriptor*,
+absl::flat_hash_map<const google::protobuf::Descriptor*,
                     CompilationEnvironments::ProcessNewEnvFn>*
     process_new_env_fns ABSL_GUARDED_BY(process_new_env_fns_mu) = nullptr;
 
@@ -128,8 +129,8 @@ CompilationEnvironments::CreateFromProto(
     const CompilationEnvironmentsProto& proto) {
   auto envs = std::make_unique<CompilationEnvironments>();
 
-  const tsl::protobuf::DescriptorPool* const pool =
-      tsl::protobuf::DescriptorPool::generated_pool();
+  const google::protobuf::DescriptorPool* const pool =
+      google::protobuf::DescriptorPool::generated_pool();
 
   for (const auto& env_proto : proto.environments()) {
     std::string fullname;
@@ -140,22 +141,21 @@ CompilationEnvironments::CreateFromProto(
                        env_proto.type_url()));
     }
 
-    const tsl::protobuf::Descriptor* const descriptor =
+    const google::protobuf::Descriptor* const descriptor =
         pool->FindMessageTypeByName(fullname);
     if (descriptor == nullptr) {
       return absl::DataLossError(absl::StrCat(
           "Unknown CompilationEnvironment message type: ", fullname));
     }
 
-    const tsl::protobuf::Message* const prototype =
-        tsl::protobuf::MessageFactory::generated_factory()->GetPrototype(
-            descriptor);
+    const google::protobuf::Message* const prototype =
+        google::protobuf::MessageFactory::generated_factory()->GetPrototype(descriptor);
     if (prototype == nullptr) {
       return absl::InternalError(absl::StrCat(
           "Unsupported CompilationEnvironment message type: ", fullname));
     }
 
-    std::unique_ptr<tsl::protobuf::Message> env(prototype->New());
+    std::unique_ptr<google::protobuf::Message> env(prototype->New());
     if (!env_proto.UnpackTo(env.get())) {
       return absl::DataLossError(absl::StrCat(
           "Unable to unpack CompilationEnvironment message of type '", fullname,
@@ -169,12 +169,11 @@ CompilationEnvironments::CreateFromProto(
 }
 
 void CompilationEnvironments::RegisterProcessNewEnvFn(
-    const tsl::protobuf::Descriptor* descriptor,
-    ProcessNewEnvFn process_new_env) {
+    const google::protobuf::Descriptor* descriptor, ProcessNewEnvFn process_new_env) {
   absl::MutexLock l(process_new_env_fns_mu);
   if (process_new_env_fns == nullptr) {
     process_new_env_fns =
-        new absl::flat_hash_map<const tsl::protobuf::Descriptor*,
+        new absl::flat_hash_map<const google::protobuf::Descriptor*,
                                 CompilationEnvironments::ProcessNewEnvFn>();
   }
   const bool inserted =
@@ -184,8 +183,21 @@ void CompilationEnvironments::RegisterProcessNewEnvFn(
                   << descriptor->full_name() << "' has already been registered";
 }
 
+void CompilationEnvironments::DeregisterProcessNewEnvFn(
+    const google::protobuf::Descriptor* descriptor) {
+  absl::MutexLock l(process_new_env_fns_mu);
+  if (process_new_env_fns == nullptr) {
+    return;
+  }
+  const auto it = process_new_env_fns->find(descriptor);
+  if (it == process_new_env_fns->end()) {
+    return;
+  }
+  process_new_env_fns->erase(it);
+}
+
 absl::Status CompilationEnvironments::InitializeAllKnownEnvs() {
-  std::vector<const tsl::protobuf::Descriptor*> descriptors;
+  std::vector<const google::protobuf::Descriptor*> descriptors;
   {
     absl::MutexLock l(process_new_env_fns_mu);
     if (process_new_env_fns == nullptr) {
@@ -208,25 +220,25 @@ absl::Status CompilationEnvironments::InitializeAllKnownEnvs() {
 }
 
 absl::Status CompilationEnvironments::AddEnv(
-    std::unique_ptr<tsl::protobuf::Message> env) {
+    std::unique_ptr<google::protobuf::Message> env) {
   if (!env) {
     return absl::InvalidArgumentError(
         "Can not add a null compilation environment.");
   }
-  const tsl::protobuf::Descriptor& descriptor = *env->GetDescriptor();
+  const google::protobuf::Descriptor& descriptor = *env->GetDescriptor();
   return AddEnvImpl(descriptor, std::move(env));
 }
 
 CompilationEnvironmentsProto CompilationEnvironments::ToProto() const {
   // Sort the environments by their message types' full names so that the
   // proto fields are deterministically ordered.
-  std::vector<const tsl::protobuf::Descriptor*> descriptors;
+  std::vector<const google::protobuf::Descriptor*> descriptors;
   descriptors.reserve(environments_.size());
   for (const auto& [descriptor, message] : environments_) {
     descriptors.push_back(descriptor);
   }
-  absl::c_sort(descriptors, [](const tsl::protobuf::Descriptor* lhs,
-                               const tsl::protobuf::Descriptor* rhs) {
+  absl::c_sort(descriptors, [](const google::protobuf::Descriptor* lhs,
+                               const google::protobuf::Descriptor* rhs) {
     return lhs->full_name() < rhs->full_name();
   });
 
@@ -239,7 +251,7 @@ CompilationEnvironmentsProto CompilationEnvironments::ToProto() const {
 
 CompilationEnvironments::ProcessNewEnvFn
 CompilationEnvironments::GetProcessNewEnvFn(
-    const tsl::protobuf::Descriptor& descriptor) {
+    const google::protobuf::Descriptor& descriptor) {
   absl::MutexLock l(process_new_env_fns_mu);
   if (process_new_env_fns == nullptr) {
     return nullptr;
@@ -262,8 +274,8 @@ void CompilationEnvironments::EnvAdded(absl::string_view env_type) {
 }
 
 absl::Status CompilationEnvironments::AddEnvImpl(
-    const tsl::protobuf::Descriptor& descriptor,
-    std::unique_ptr<tsl::protobuf::Message> env) {
+    const google::protobuf::Descriptor& descriptor,
+    std::unique_ptr<google::protobuf::Message> env) {
   // Check if we already have an environment of env's type
   if (environments_.contains(&descriptor)) {
     return absl::AlreadyExistsError(absl::StrCat(
@@ -276,16 +288,16 @@ absl::Status CompilationEnvironments::AddEnvImpl(
     return absl::InvalidArgumentError(absl::StrCat(
         "Unknown CompilationEnvironment type ", descriptor.full_name()));
   }
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<tsl::protobuf::Message> processed_env,
+  TF_ASSIGN_OR_RETURN(std::unique_ptr<google::protobuf::Message> processed_env,
                       process_new_env(std::move(env)));
 
   // Check for unknown fields
-  const tsl::protobuf::UnknownFieldSet& unknown_fields =
+  const google::protobuf::UnknownFieldSet& unknown_fields =
       processed_env->GetReflection()->GetUnknownFields(*processed_env);
   std::vector<int> unknown_tags;
   unknown_tags.reserve(unknown_fields.field_count());
   for (int i = 0; i < unknown_fields.field_count(); ++i) {
-    const tsl::protobuf::UnknownField& field = unknown_fields.field(i);
+    const google::protobuf::UnknownField& field = unknown_fields.field(i);
     unknown_tags.push_back(field.number());
   }
   if (!unknown_tags.empty()) {

xla/service/compilation_environments.h

@@ -24,8 +24,8 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "google/protobuf/descriptor.h"
 #include "xla/xla.pb.h"
-#include "tsl/platform/protobuf.h"
 
 namespace xla {
 
@@ -34,7 +34,7 @@ namespace xla {
 //
 // CompilationEnvironments uses lazy initialization, (see GetEnv() for more
 // details). Lazy initialization is used so we can avoid:
-// A) Requiring every code path to explitily construct all needed compilation
+// A) Requiring every code path to explicitly construct all needed compilation
 //    environments, particularly when the default constructed environment is
 //    all we need AND
 // B) Requiring CompilationEnvironments to implicitly construct all needed
@@ -45,8 +45,8 @@ namespace xla {
 class CompilationEnvironments {
  public:
   using ProcessNewEnvFn =
-      std::function<absl::StatusOr<std::unique_ptr<tsl::protobuf::Message>>(
-          std::unique_ptr<tsl::protobuf::Message>)>;
+      std::function<absl::StatusOr<std::unique_ptr<google::protobuf::Message>>(
+          std::unique_ptr<google::protobuf::Message>)>;
 
   CompilationEnvironments() = default;
   CompilationEnvironments(const CompilationEnvironments& rhs) { *this = rhs; }
@@ -72,17 +72,22 @@ class CompilationEnvironments {
   //
   // REQUIRES:
   // - The output is *not* allowed to be null, even for null input.
-  static void RegisterProcessNewEnvFn(
-      const tsl::protobuf::Descriptor* descriptor,
-      ProcessNewEnvFn process_new_env);
+  // - `descriptor` must stay alive until the process ends, or until
+  //   `DeregisterProcessNewEnvFn` is called.
+  static void RegisterProcessNewEnvFn(const google::protobuf::Descriptor* descriptor,
+                                      ProcessNewEnvFn process_new_env);
+
+  // Deregisters the ProcessNewEnvFn for the given proto descriptor, if one
+  // exists.
+  static void DeregisterProcessNewEnvFn(const google::protobuf::Descriptor* descriptor);
 
   // Adds env to the list of CompilationEnvironments. If an environment with
   // the same proto descriptor has already been added, returns an error.
   //
   // All added environments are processed via registered ProcessNewEnvFns. If
   // such a function was not regitered for env's proto descriptor or env's
   // proto type is unknown, an error will be returned.
-  absl::Status AddEnv(std::unique_ptr<tsl::protobuf::Message> env);
+  absl::Status AddEnv(std::unique_ptr<google::protobuf::Message> env);
 
   // Returns the CompilationEnvironment corresponding to T. If such an
   // environment has not been added, ProcessNewEnvFn(nullptr) will be added and
@@ -116,7 +121,7 @@ class CompilationEnvironments {
   // Returns the ProcessNewEnvFn for the given env type. Returns nullptr if no
   // ProcessNewEnvFn has been registered for the env type.
   static ProcessNewEnvFn GetProcessNewEnvFn(
-      const tsl::protobuf::Descriptor& descriptor);
+      const google::protobuf::Descriptor& descriptor);
 
   // Called by GetEnv(), when it lazily creates a new environment, to globally
   // track stats about how many such environments are created by
@@ -128,11 +133,11 @@ class CompilationEnvironments {
   // are added to CompilationEnvironments.
   static void EnvAdded(absl::string_view env_type);
 
-  absl::Status AddEnvImpl(const tsl::protobuf::Descriptor& descriptor,
-                          std::unique_ptr<tsl::protobuf::Message> env);
+  absl::Status AddEnvImpl(const google::protobuf::Descriptor& descriptor,
+                          std::unique_ptr<google::protobuf::Message> env);
 
-  absl::flat_hash_map<const tsl::protobuf::Descriptor*,
-                      std::unique_ptr<tsl::protobuf::Message>>
+  absl::flat_hash_map<const google::protobuf::Descriptor*,
+                      std::unique_ptr<google::protobuf::Message>>
       environments_;
 };
 
@@ -158,7 +163,7 @@ T& CompilationEnvironments::GetMutableEnv() {
     it = environments_.find(descriptor);
   }
 
-  return tsl::protobuf::DownCastToGenerated<T>(*it->second);
+  return google::protobuf::DownCastToGenerated<T>(*it->second);
 }
 
 template <typename T>