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

[tf.data] Maintain a reference on the FunctionBufferingResource while a get-next operation is active.

Browse Source 
Previously, the reference count on a FunctionBufferingResource could drop to 0 and it could be deleted (e.g. by a DestroyResourceOp) while a get-next operation is active on it. This would lead to use-after-free errors.

PiperOrigin-RevId: 190878208
This commit is contained in:
Derek Murray 2018-03-28 20:44:51 -07:00 committed by TensorFlower Gardener
parent a5a90e6b55
commit 789e442513
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/contrib/data/kernels/prefetching_kernels.cc
View File

@ -374,25 +374,27 @@ class FunctionBufferingResourceGetNextOp : public AsyncOpKernel {
OP_REQUIRES_OK_ASYNC(
ctx, LookupResource<FunctionBufferingResource>(ctx, handle, &buffer),
done);
core::ScopedUnref s(buffer);
if (buffer->Finished()) {
buffer->Unref();
ctx->SetStatus(errors::OutOfRange("end_of_sequence"));
done();
return;
}
FunctionBufferCallback callback =
[ctx, done](const BufferElement& buffer_element) {
[ctx, buffer, done](const BufferElement& buffer_element) {
Status s = buffer_element.status;
if (!s.ok()) {
ctx->SetStatus(s);
buffer->Unref();
done();
return;
}
for (size_t i = 0; i < buffer_element.value.size(); ++i) {
ctx->set_output(i, buffer_element.value[i]);
}
buffer->Unref();
done();
};
buffer->MaybeGet(std::move(callback));