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

Deregister cancellation callback before invoking pros_done or cons_done

Browse Source 
If we invoke done() in the consumer or provider before deregistering the
cancellation callback, the cancellation manager may already be gone.

PiperOrigin-RevId: 338358532
Change-Id: I9dff626dcf25734c5888f0003a4e21ba0d89274f
This commit is contained in:
Ran Chen 2020-10-21 16:01:04 -07:00 committed by TensorFlower Gardener
parent 8f8fffcd6e
commit 01562ab770
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tensorflow/core/common_runtime/buf_rendezvous.cc
View File

@ -25,6 +25,15 @@ limitations under the License.
#include "tensorflow/core/lib/core/notification.h" #include "tensorflow/core/lib/core/notification.h"
namespace tensorflow { namespace tensorflow {
namespace {
void DeregisterCancellation(BufRendezvous::Hook* h) {
if (h->cancellation_manager != nullptr) {
h->cancellation_manager->DeregisterCallback(h->cancellation_token);
h->cancellation_manager = nullptr;
h->cancellation_token = CancellationManager::kInvalidToken;
}
}
} // namespace
BufRendezvous::~BufRendezvous() { BufRendezvous::~BufRendezvous() {
mutex_lock l(mu_); mutex_lock l(mu_);
@ -129,6 +138,7 @@ void BufRendezvous::ProvideBuf(const string& key, Device* dev,
} }
} while (false); } while (false);
if (h) { if (h) {
DeregisterCancellation(h);
h->cons_cb(Status::OK(), h); h->cons_cb(Status::OK(), h);
} }
if (!providebuf_status.ok()) { if (!providebuf_status.ok()) {
@ -197,6 +207,7 @@ void BufRendezvous::ConsumeBuf(const string& key, const string& device_name,
} }
} while (false); } while (false);
if (existing_hook) { if (existing_hook) {
DeregisterCancellation(existing_hook);
existing_hook->cons_cb(Status::OK(), existing_hook); existing_hook->cons_cb(Status::OK(), existing_hook);
return; return;
} }
@ -230,9 +241,6 @@ void BufRendezvous::CancelHook(const string& key) {
/*static*/ /*static*/
void BufRendezvous::DoneWithHook(Hook* h) { void BufRendezvous::DoneWithHook(Hook* h) {
if (h->cancellation_manager != nullptr) {
h->cancellation_manager->DeregisterCallback(h->cancellation_token);
}
h->prod_cb(Status::OK()); h->prod_cb(Status::OK());
delete h; delete h;
} }