Fix a race condition in tensorflow::ProfilerServer.

server_ is accessed concurrently from two threads without any locking. Amongst other things this could lead to a deadlock if a server was started and then immediately destroyed, because the
if (server_) ...
condition in ~ProfilerServer could run before server_ had been assigned in the worker thread.

Rather than add locking, a simpler solution is just to remove the thread altogether. There is no need to use a thread here: we simply need to call Wait() in the destructor, because gRPC has its own threading internally.

PiperOrigin-RevId: 294996411
Change-Id: Idd507d1ea97642e540fdec520900a3c491d34e5a

tensorflow/core/profiler/rpc/profiler_server.cc

@@ -29,21 +29,14 @@ limitations under the License.
 namespace tensorflow {
 
 void ProfilerServer::StartProfilerServer(int32 port) {
-  Env* env = Env::Default();
-  auto start_server = [port, this]() {
-    string server_address = absl::StrCat("0.0.0.0:", port);
-    std::unique_ptr<grpc::ProfilerService::Service> service =
-        CreateProfilerService();
-    ::grpc::ServerBuilder builder;
-    builder.AddListeningPort(server_address,
-                             ::grpc::InsecureServerCredentials());
-    builder.RegisterService(service.get());
-    server_ = builder.BuildAndStart();
-    LOG(INFO) << "Profiling Server listening on " << server_address;
-    server_->Wait();
-  };
-  server_thread_ =
-      WrapUnique(env->StartThread({}, "ProfilerServer", start_server));
+  string server_address = absl::StrCat("0.0.0.0:", port);
+  std::unique_ptr<grpc::ProfilerService::Service> service =
+      CreateProfilerService();
+  ::grpc::ServerBuilder builder;
+  builder.AddListeningPort(server_address, ::grpc::InsecureServerCredentials());
+  builder.RegisterService(service.get());
+  server_ = builder.BuildAndStart();
+  LOG(INFO) << "Profiling Server listening on " << server_address;
 }
 
 void ProfilerServer::MaybeStartProfilerServer() {
@@ -67,7 +60,10 @@ void ProfilerServer::MaybeStartProfilerServer() {
 }
 
 ProfilerServer::~ProfilerServer() {
-  if (server_) server_->Shutdown();
+  if (server_) {
+    server_->Shutdown();
+    server_->Wait();
+  }
 }
 
 }  // namespace tensorflow

tensorflow/core/profiler/rpc/profiler_server.h

@@ -35,7 +35,6 @@ class ProfilerServer {
 
  private:
   std::unique_ptr<::grpc::Server> server_;
-  std::unique_ptr<Thread> server_thread_;
 };
 
 }  // namespace tensorflow