[libtpu] Make copies of data returned by getenv.

According to http://www.cplusplus.com/reference/cstdlib/getenv/:
The pointer returned points to an internal memory block, whose content
or validity may be altered by further calls to getenv (but not by
other library functions).

This appears to have been working without the defensive copies, but we
should follow the spec just in case.

This change also refactors TfTpu_Initialize to not use the argc/argv
convention for simplicity (the implementation doesn't actually require
a nullptr at the end of argv).

PiperOrigin-RevId: 354396268
Change-Id: Ie14f13d4a67e1012dc13af63df6a9a4fcb22feeb

tensorflow/core/tpu/libtftpu.h

@@ -41,8 +41,8 @@ limitations under the License.
 extern "C" {
 #endif
 
-TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library, int argc,
-                                        const char** argv);
+TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library, int num_args,
+                                        const char** args);
 
 #ifdef __cplusplus
 }

tensorflow/core/tpu/tpu_api_dlsym_initializer.cc

@@ -45,17 +45,18 @@ Status InitializeTpuLibrary(void* library_handle) {
   Status s = InitializeTpuStructFns(library_handle);
 
   // Retrieve arguments from environment if applicable
-  std::vector<const char*> argv_ptr = GetLibTpuInitArguments();
+  std::pair<std::vector<std::string>, std::vector<const char*> > args =
+      GetLibTpuInitArguments();
 
   // TPU platform registration must only be performed after the library is
   // loaded. We do not want to register a TPU platform in XLA without the
   // supporting library providing the necessary APIs.
   if (s.ok()) {
-    void (*initialize_fn)(bool init_library, int argc, const char** argv);
+    void (*initialize_fn)(bool init_library, int num_args, const char** args);
     initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
         dlsym(library_handle, "TfTpu_Initialize"));
-    (*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1,
-                     /*argv=*/argv_ptr.data());
+    (*initialize_fn)(/*init_library=*/true, args.second.size(),
+                     args.second.data());
 
     RegisterTpuPlatform();
   }

tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc

@@ -42,17 +42,18 @@ Status InitializeTpuLibrary(void* library_handle) {
   Status s = SetExecutorStructFn(library_handle);
 
   // Retrieve arguments from environment if applicable
-  std::vector<const char*> argv_ptr = GetLibTpuInitArguments();
+  std::pair<std::vector<std::string>, std::vector<const char*> > args =
+      GetLibTpuInitArguments();
 
   // TPU platform registration must only be performed after the library is
   // loaded. We do not want to register a TPU platform in XLA without the
   // supporting library providing the necessary APIs.
   if (s.ok()) {
-    void (*initialize_fn)(bool init_library, int argc, const char** argv);
+    void (*initialize_fn)(bool init_library, int num_args, const char** args);
     initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
         dlsym(library_handle, "TfTpu_Initialize"));
-    (*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1,
-                     /*argv=*/argv_ptr.data());
+    (*initialize_fn)(/*init_library=*/true, args.second.size(),
+                     args.second.data());
 
     RegisterTpuPlatform();
   }

tensorflow/core/tpu/tpu_initializer_helper.cc

@@ -22,17 +22,17 @@ limitations under the License.
 namespace tensorflow {
 namespace tpu {
 
-std::vector<const char*> GetLibTpuInitArguments() {
+std::pair<std::vector<std::string>, std::vector<const char*>>
+GetLibTpuInitArguments() {
+  // We make copies of the arguments returned by getenv because the memory
+  // returned may be altered or invalidated by further calls to getenv.
+  std::vector<std::string> argv;
   std::vector<const char*> argv_ptr;
-  std::vector<absl::string_view> argv;
 
-  // Retrieve arguments from environment if applicable
+  // Retrieve arguments from environment if applicable.
   char* env = getenv("LIBTPU_INIT_ARGS");
   if (env != nullptr) {
     // TODO(frankchn): Handles quotes properly if necessary.
-    // env pointer is already pointing to an allocated memory block.
-    // absl::StrSplit returns a string_view that returns a vector of pointers
-    // into that memory block. This means that we don't need to manage memory.
     argv = absl::StrSplit(env, ' ');
   }
 
@@ -42,7 +42,7 @@ std::vector<const char*> GetLibTpuInitArguments() {
   }
   argv_ptr.push_back(nullptr);
 
-  return argv_ptr;
+  return {argv, argv_ptr};
 }
 
 }  // namespace tpu

tensorflow/core/tpu/tpu_initializer_helper.h

@@ -16,14 +16,17 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_
 #define TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_
 
+#include <string>
 #include <vector>
 
 namespace tensorflow {
 namespace tpu {
 
-// This returns an extra nullptr at the end (per the C standard), but this
-// should not be counted for 'argc'.
-std::vector<const char*> GetLibTpuInitArguments();
+// Returns arguments (e.g. flags) set in the LIBTPU_INIT_ARGS environment
+// variable. The first return value is the arguments, the second return value is
+// pointers to the arguments suitable for passing into the C API.
+std::pair<std::vector<std::string>, std::vector<const char*>>
+GetLibTpuInitArguments();
 
 }  // namespace tpu
 }  // namespace tensorflow