1 files changed, 16 insertions, 22 deletions

diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc
index 8b9b3da21c..778cb667e2 100644
--- a/tensorflow/c/c_api.cc
+++ b/tensorflow/c/c_api.cc
@@ -63,6 +63,7 @@ limitations under the License.
 // brain namespace because we are defining 'extern "C"' functions.
 using tensorflow::AllocationDescription;
 using tensorflow::DataType;
+using tensorflow::ExtendSessionGraphHelper;
 using tensorflow::Graph;
 using tensorflow::GraphDef;
 using tensorflow::mutex_lock;
@@ -640,11 +641,11 @@ Status MessageToBuffer(const tensorflow::protobuf::Message& in,
 }
 
 void RecordMutation(TF_Graph* graph, const TF_Operation& op,
-                    const char* mutation_type)
-    EXCLUSIVE_LOCKS_REQUIRED(graph->mu) {
+                    const char* mutation_type) {
   // If any session has already run this node_id, mark this session as
   // unrunnable.
   for (auto it : graph->sessions) {
+    mutex_lock session_lock(it.first->mu);
     if (it.first->last_num_graph_nodes > op.node.id()) {
       it.second = FailedPrecondition(
           "Operation '", op.node.DebugString(), "' was changed by ",
@@ -713,10 +714,12 @@ Status LoadLibrary(const char* library_filename, void** result,
 // TODO(josh11b,mrry): Change Session to be able to use a Graph*
 // directly, instead of requiring us to serialize to a GraphDef and
 // call Session::Extend().
-bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status)
-    EXCLUSIVE_LOCKS_REQUIRED(session->mu) {
+bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status) {
   if (session->graph != nullptr) {
+    // Take the graph lock before the session lock to avoid deadlock. This is
+    // safe since session->graph does not change.
     session->graph->mu.lock();
+    mutex_lock session_lock(session->mu);
     const Graph& graph = session->graph->graph;
 
     status->status = session->graph->sessions[session];
@@ -2571,12 +2574,9 @@ void TF_SessionRun(TF_Session* session, const TF_Buffer* run_options,
   // TODO(josh11b,mrry): Change Session to be able to use a Graph*
   // directly, instead of requiring us to serialize to a GraphDef and
   // call Session::Extend().
-  {
-    mutex_lock l(session->mu);
-    if (session->extend_before_run &&
-        !tensorflow::ExtendSessionGraphHelper(session, status)) {
-      return;
-    }
+  if (session->extend_before_run &&
+      !ExtendSessionGraphHelper(session, status)) {
+    return;
   }
 
   TF_Run_Setup(noutputs, output_values, status);
@@ -2612,12 +2612,9 @@ void TF_SessionPRunSetup(TF_Session* session, const TF_Output* inputs,
                          const char** handle, TF_Status* status) {
   *handle = nullptr;
 
-  {
-    mutex_lock l(session->mu);
-    if (session->extend_before_run &&
-        !tensorflow::ExtendSessionGraphHelper(session, status)) {
-      return;
-    }
+  if (session->extend_before_run &&
+      !ExtendSessionGraphHelper(session, status)) {
+    return;
   }
 
   std::vector<string> input_names(ninputs);
@@ -2659,12 +2656,9 @@ void TF_SessionPRun(TF_Session* session, const char* handle,
   // TODO(josh11b,mrry): Change Session to be able to use a Graph*
   // directly, instead of requiring us to serialize to a GraphDef and
   // call Session::Extend().
-  {
-    mutex_lock l(session->mu);
-    if (session->extend_before_run &&
-        !tensorflow::ExtendSessionGraphHelper(session, status)) {
-      return;
-    }
+  if (session->extend_before_run &&
+      !ExtendSessionGraphHelper(session, status)) {
+    return;
   }
 
   TF_Run_Setup(noutputs, output_values, status);