From fb1069781ffcbac222392a68c01a45fae264888e Mon Sep 17 00:00:00 2001 From: Brennan Saeta Date: Sat, 28 Apr 2018 10:51:32 -0700 Subject: [PATCH] [tf.data] Use core::ScopedUnref to avoid resource leakage. If for whatever reason iterator_resource->set_iterator did not return Status::OK(), we would leak a reference on the iterator_resource. With this change, we won't leak the resource. PiperOrigin-RevId: 194662412 --- tensorflow/core/kernels/data/iterator_ops.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc index f5db97f..a2f6c5f 100644 --- a/tensorflow/core/kernels/data/iterator_ops.cc +++ b/tensorflow/core/kernels/data/iterator_ops.cc @@ -584,9 +584,9 @@ class MakeIteratorOp : public OpKernel { IteratorResource* iterator_resource; OP_REQUIRES_OK( ctx, LookupResource(ctx, HandleFromInput(ctx, 1), &iterator_resource)); + core::ScopedUnref unref(iterator_resource); OP_REQUIRES_OK(ctx, iterator_resource->set_iterator( dataset->MakeIterator("Iterator"))); - iterator_resource->Unref(); } }; -- 2.7.4