Fixing race condition at Module::forward method (#21398)

Yuri Putivsky · facebook-github-bot · commit 42b2f56124f4 · 2019-06-04T23:03:25.000-07:00
Summary: Pull Request resolved: #21398 Module::forward method calls find_method() function potentially in multiple threads. Internally it calls find_offset() method and reads dict_ object. If the correspondent name is not in a dictionary thread call insert() method and modifies dict_ object. At the same time when first thread modifies dict_ object another thread can enter forward()->find_method()->find_offset() path and access dict_ object for reading while it have been modified -> crash. Moved mutex protection up to protect both calls find_offset() and insert(). Consider to use C++ 17 shared_mutex locking object instead of recursive_mutex object. Reviewed By: bddppq Differential Revision: D15638942 fbshipit-source-id: ca6a453448302a0b3666c87724755fa4e9ce242f
diff --git a/torch/csrc/jit/script/module.h b/torch/csrc/jit/script/module.h
@@ -304,6 +304,12 @@ struct TORCH_API Module {
     return offset ? modules_[*offset] : nullptr;
   }
   Method* find_method(const std::string& name) const {
+    // find_offset() method reads "dict_" object.
+    // Lock because another thread can modify "dict_" object at the same time
+    // calling insert() method.
+    // Ideally recursive_mutex should be replaced with shared_mutex (C++ 17)
+    // for the performance reasons.
+    std::unique_lock<std::recursive_mutex> keeper(create_method_guard_);
     auto offset = find_offset(name, EntityType::METHOD);
     if (offset) {
       return methods_[*offset].get();
@@ -314,7 +320,6 @@ struct TORCH_API Module {
       // but we have to update the internal Method cache.
       // This can be removed when class_compilation_unit() is the source of
       // truth for methods.
-      std::lock_guard<std::recursive_mutex> guard(create_method_guard_);
       Module* mutable_this = const_cast<Module*>(this);
       std::unique_ptr<Method> m(new Method(mutable_this, fn));
       return mutable_this
