[Runtime] Fix shared resource_mgr double free corruption bug. (#286)

Author: shanshanpt

tensorflow/core/common_runtime/device.cc

@@ -28,6 +28,7 @@ Device::Device(Env* env, const DeviceAttributes& device_attributes)
   CHECK(DeviceNameUtils::ParseFullName(name(), &parsed_name_))
       << "Invalid device name: " << name();
   rmgr_ = new ResourceMgr(parsed_name_.job);
+  owned_rmgr_ = true;
 }
 
 Device::Device(Env* env, const DeviceAttributes& device_attributes,
@@ -41,8 +42,10 @@ Device::Device(Env* env, const DeviceAttributes& device_attributes,
           dev_rmgr_map->device_rmgr_map.end()) {
     rmgr_ = const_cast<DeviceResourceMgrMap*>(dev_rmgr_map)->device_rmgr_map[name()];
     LOG(INFO) << "Device " << name() << " got a shared resource_mgr: " << rmgr_;
+    owned_rmgr_ = false;
   } else {
     rmgr_ = new ResourceMgr(parsed_name_.job);
+    owned_rmgr_ = true;
   }
 }
 

tensorflow/core/common_runtime/device.h

@@ -191,7 +191,9 @@ class Device : public DeviceBase {
 
  protected:
   void DeleteResourceMgr() {
-    delete rmgr_;
+    if (owned_rmgr_) {
+      delete rmgr_;
+    }
     rmgr_ = nullptr;
   }
 
@@ -204,6 +206,7 @@ class Device : public DeviceBase {
 
   // Resources associated w/ this device. E.g., shared variables, etc.
   ResourceMgr* rmgr_ = nullptr;
+  bool owned_rmgr_ = true;
 
   TF_DISALLOW_COPY_AND_ASSIGN(Device);
 };

tensorflow/core/common_runtime/direct_session.cc

@@ -286,7 +286,7 @@ class DirectSessionFactory : public SessionFactory {
 
     DeviceMgr* device_mgr = new DeviceMgr(std::move(devices));
 
-    SessionGroup* session_group = new SessionGroup();
+    SessionGroup* session_group = new SessionGroup(shared_rmgr);
 #ifdef TENSORFLOW_USE_NUMA
     DirectSession* leader_session =
         new DirectSession(options, device_mgr, true, this,

tensorflow/core/public/session.h

@@ -30,6 +30,7 @@ limitations under the License.
 
 namespace tensorflow {
 class DeviceMgr;
+class ResourceMgr;
 
 namespace thread {
 
@@ -270,7 +271,12 @@ class Session {
 
 class SessionGroup {
  public:
+  SessionGroup() : shared_resource_mgr_(nullptr) {}
+  SessionGroup(ResourceMgr* mgr) : shared_resource_mgr_(mgr) {}
   ~SessionGroup() {
+    if (shared_resource_mgr_) {
+      delete shared_resource_mgr_;
+    }
   }
 
   Status Close() {
@@ -375,6 +381,7 @@ class SessionGroup {
   std::vector<std::unique_ptr<Session>> sessions_;
   int32_t session_num_ = 0;
   std::atomic<int64_t> serving_index_{0};
+  ResourceMgr* shared_resource_mgr_ = nullptr;
 
   Status GetServingSessionId(int32_t* serving_id, int32_t hint_id = -1) {
     if (session_num_ < 1) {