From 789e442513e85ab1caeb1e03997b0aafa3cd76d7 Mon Sep 17 00:00:00 2001 From: Derek Murray Date: Wed, 28 Mar 2018 20:44:51 -0700 Subject: [PATCH] [tf.data] Maintain a reference on the FunctionBufferingResource while a get-next operation is active. 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 --- tensorflow/contrib/data/kernels/prefetching_kernels.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/contrib/data/kernels/prefetching_kernels.cc b/tensorflow/contrib/data/kernels/prefetching_kernels.cc index f51570d..2afb8db 100644 --- a/tensorflow/contrib/data/kernels/prefetching_kernels.cc +++ b/tensorflow/contrib/data/kernels/prefetching_kernels.cc @@ -374,25 +374,27 @@ class FunctionBufferingResourceGetNextOp : public AsyncOpKernel { OP_REQUIRES_OK_ASYNC( ctx, LookupResource(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)); -- 2.7.4