From a85beef408a3f5dcd93158e88a0aade6cb70f8c5 Mon Sep 17 00:00:00 2001 From: Kuangyuan Chen Date: Wed, 8 Jul 2020 09:56:54 -0700 Subject: [PATCH] Fix a race condition in attr_builder where the AttrTypeMap for the same op_name might be created multiple times, causing memory leak. PiperOrigin-RevId: 320208354 Change-Id: I36b7f077ae6645160198c55d0441d482141aa707 --- tensorflow/core/common_runtime/eager/attr_builder.cc | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/common_runtime/eager/attr_builder.cc b/tensorflow/core/common_runtime/eager/attr_builder.cc index 2d85b2f764d..fd79a82e4b2 100644 --- a/tensorflow/core/common_runtime/eager/attr_builder.cc +++ b/tensorflow/core/common_runtime/eager/attr_builder.cc @@ -71,7 +71,15 @@ Status AttrTypeMapForOp(const char* op_name, const AttrTypeMap** out, *out = gtl::FindPtrOrNull(*OpNameToAttrTypeMap(), op_name); if (*out != nullptr) return Status::OK(); } + mutex_lock l(g_op_name_to_attr_type_map_lock); + + // Check the existence of AttrTypeMap for op_name again because another thread + // may insert this map after the tf_shared_lock is released but before the + // mutex_lock is acquired. + *out = gtl::FindPtrOrNull(*OpNameToAttrTypeMap(), op_name); + if (*out != nullptr) return Status::OK(); + const OpDef* op_def = nullptr; Status s = OpDefForOp(op_name, &op_def); if (errors::IsNotFound(s)) { @@ -121,7 +129,9 @@ Status AttrTypeMapForOp(const char* op_name, const AttrTypeMap** out, gtl::InsertIfNotPresent(m.get(), attr.name(), t); } *out = m.get(); - (*OpNameToAttrTypeMap())[op_name] = m.release(); + auto r = OpNameToAttrTypeMap()->emplace(op_name, m.release()); + DCHECK(r.second) << "AttrTypeMap already exists for " << op_name; + return Status::OK(); }