experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix use-after-free of CancellationManager in LocalRendezvousImpl.

Browse Source 
Previously, we were invoking the CancellationManager in ~Item, which runs after the done callback. However, the CancellationManager was borrowed from the calling RecvOp, and it will tend to be deleted synchronously when the done callback executes.

PiperOrigin-RevId: 266567112
This commit is contained in:
Derek Murray 2019-08-31 11:13:30 -07:00
parent 25006be096
commit e8ee6b8d67
1 changed files with 18 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
tensorflow/core/framework/rendezvous.cc
View File

@ -264,7 +264,24 @@ class LocalRendezvousImpl : public Rendezvous {
VLOG(2) << "Enqueue Recv Item (key:" << key.FullKey() << "). "; VLOG(2) << "Enqueue Recv Item (key:" << key.FullKey() << "). ";
Item* item = new Item; Item* item = new Item;
item->waiter = std::move(done);
if (cm != nullptr) {
auto wrapped_done = std::bind(
[cm, token](const DoneCallback& done,
// Begin unbound arguments.
const Status& s, const Args& send_args,
const Args& recv_args, const Tensor& v, bool dead) {
cm->TryDeregisterCallback(token);
done(s, send_args, recv_args, v, dead);
},
std::move(done), std::placeholders::_1, std::placeholders::_2,
std::placeholders::_3, std::placeholders::_4,
std::placeholders::_5);
item->waiter = std::move(wrapped_done);
} else {
item->waiter = std::move(done);
}
item->recv_args = recv_args; item->recv_args = recv_args;
item->cancellation_token = token; item->cancellation_token = token;
if (item->recv_args.device_context) { if (item->recv_args.device_context) {
@ -332,11 +349,6 @@ class LocalRendezvousImpl : public Rendezvous {
if (recv_args.device_context) { if (recv_args.device_context) {
recv_args.device_context->Unref(); recv_args.device_context->Unref();
} }
auto* cm = recv_args.cancellation_manager;
if (cancellation_token != CancellationManager::kInvalidToken &&
cm != nullptr) {
cm->TryDeregisterCallback(cancellation_token);
}
} }
// Returns true iff this item represents a value being sent. // Returns true iff this item represents a value being sent.