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
This commit is contained in:
Kuangyuan Chen 2020-07-08 09:56:54 -07:00 committed by TensorFlower Gardener
parent b1cdfe3b9b
commit a85beef408

View File

@ -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();
}