diff options
| author |  Derek Murray <mrry@google.com> | 2018-03-28 14:30:39 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-03-28 14:33:15 -0700 |
| commit | ef6552b544b3c3bf6808be807b30dd9bd4f19669 (patch) | |
| tree | 1bbc445e961434c666ae00c1be871914f2948b24 | |
| parent | d355f4e2644b68ea643f573c564936ec23b93787 (diff) | |
[tf.data] Fix reference leak in FunctionBufferingResource.
Previously, the FunctionBufferingResource's destructor would never be called, which led to use-after-free (of the underlying Device object) errors in the prefetching function.
PiperOrigin-RevId: 190834415
| -rw-r--r-- | tensorflow/contrib/cmake/tf_tests.cmake | 1 | ||||
| -rw-r--r-- | tensorflow/contrib/data/kernels/prefetching_kernels.cc | 1 |
2 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/contrib/cmake/tf_tests.cmake b/tensorflow/contrib/cmake/tf_tests.cmake index 237f4fe33a..f793877c8b 100644 --- a/tensorflow/contrib/cmake/tf_tests.cmake +++ b/tensorflow/contrib/cmake/tf_tests.cmake @@ -281,6 +281,7 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${tensorflow_source_dir}/tensorflow/python/data/kernel_tests/iterator_ops_cluster_test.py" "${tensorflow_source_dir}/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py" # Deadlocks "${tensorflow_source_dir}/tensorflow/contrib/data/python/kernel_tests/sloppy_transformation_dataset_op_test.py" # b/65430561 + "${tensorflow_source_dir}/tensorflow/contrib/data/python/kernel_tests/prefetching_ops_test.py" # Segfaults on Windows. # tensor_forest tests (also note that we exclude the hybrid tests for now) "${tensorflow_source_dir}/tensorflow/contrib/tensor_forest/python/kernel_tests/count_extremely_random_stats_op_test.py" # Results in wrong order. "${tensorflow_source_dir}/tensorflow/contrib/tensor_forest/python/kernel_tests/sample_inputs_op_test.py" # Results in wrong order. diff --git a/tensorflow/contrib/data/kernels/prefetching_kernels.cc b/tensorflow/contrib/data/kernels/prefetching_kernels.cc index 79d1fc3494..f51570db85 100644 --- a/tensorflow/contrib/data/kernels/prefetching_kernels.cc +++ b/tensorflow/contrib/data/kernels/prefetching_kernels.cc @@ -314,6 +314,7 @@ class FunctionBufferResourceHandleOp : public OpKernel { source_device, target_device, func_args, thread_pool_size_); return Status::OK(); })); + core::ScopedUnref s(buffer); OP_REQUIRES_OK(ctx, buffer->Instantiate()); initialized_ = true; } |