Avoid unnecessary buffer allocations & deallocations

Before this change, when we reached the end of a file, we would
 (1) clear the existing buffer (which at large buffer sizes typically involved
     deallocating it).
 (2) reserve a buffer (which at large buffer sizes is non-trivial)
 (3) realize we had reached EoF, and therefore clear the buffer, deallocating
     it again.

With this change, whenever the buffered reader detects an EoF condition, we
remember it, so that we can short-circuit the above logic.

The above optimization results in a more than 25x performance improvement for
large buffers reading small files.

PiperOrigin-RevId: 167766751

3 files changed, 62 insertions, 0 deletions

diff --git a/tensorflow/core/lib/io/buffered_inputstream.cc b/tensorflow/core/lib/io/buffered_inputstream.cc
index 6f72da4713..b247e9c575 100644
--- a/tensorflow/core/lib/io/buffered_inputstream.cc
+++ b/tensorflow/core/lib/io/buffered_inputstream.cc
@@ -41,9 +41,18 @@ BufferedInputStream::~BufferedInputStream() {
 }
 
 Status BufferedInputStream::FillBuffer() {
+  if (!file_status_.ok()) {
+    pos_ = 0;
+    limit_ = 0;
+    return file_status_;
+  }
   Status s = input_stream_->ReadNBytes(size_, &buf_);
   pos_ = 0;
   limit_ = buf_.size();
+  if (buf_.empty()) {
+    DCHECK(!s.ok());
+    file_status_ = s;
+  }
   return s;
 }
 
@@ -82,6 +91,9 @@ Status BufferedInputStream::ReadNBytes(int64 bytes_to_read, string* result) {
                                    bytes_to_read);
   }
   result->clear();
+  if (!file_status_.ok() && bytes_to_read > 0) {
+    return file_status_;
+  }
   result->reserve(bytes_to_read);
 
   Status s;
@@ -91,6 +103,8 @@ Status BufferedInputStream::ReadNBytes(int64 bytes_to_read, string* result) {
       s = FillBuffer();
       // If we didn't read any bytes, we're at the end of the file; break out.
       if (limit_ == 0) {
+        DCHECK(!s.ok());
+        file_status_ = s;
         break;
       }
     }
@@ -124,6 +138,9 @@ Status BufferedInputStream::SkipNBytes(int64 bytes_to_skip) {
     Status s = input_stream_->SkipNBytes(bytes_to_skip - (limit_ - pos_));
     pos_ = 0;
     limit_ = 0;
+    if (errors::IsOutOfRange(s)) {
+      file_status_ = s;
+    }
     return s;
   }
   return Status::OK();
@@ -163,6 +180,7 @@ Status BufferedInputStream::ReadAll(string* result) {
   }
 
   if (errors::IsOutOfRange(status)) {
+    file_status_ = status;
     return Status::OK();
   }
   return status;
@@ -172,6 +190,7 @@ Status BufferedInputStream::Reset() {
   TF_RETURN_IF_ERROR(input_stream_->Reset());
   pos_ = 0;
   limit_ = 0;
+  file_status_ = Status::OK();
   return Status::OK();
 }
 
diff --git a/tensorflow/core/lib/io/buffered_inputstream.h b/tensorflow/core/lib/io/buffered_inputstream.h
index b37766005a..2b824f35f8 100644
--- a/tensorflow/core/lib/io/buffered_inputstream.h
+++ b/tensorflow/core/lib/io/buffered_inputstream.h
@@ -94,6 +94,9 @@ class BufferedInputStream : public InputStreamInterface {
   size_t pos_ = 0;    // current position in buf_.
   size_t limit_ = 0;  // just past the end of valid data in buf_.
   bool owns_input_stream_ = false;
+  // When EoF is reached, file_status_ contains the status to skip unnecessary
+  // buffer allocations.
+  Status file_status_ = Status::OK();
 
   TF_DISALLOW_COPY_AND_ASSIGN(BufferedInputStream);
 };
diff --git a/tensorflow/core/lib/io/buffered_inputstream_test.cc b/tensorflow/core/lib/io/buffered_inputstream_test.cc
index 7265101e1b..49b2b1a861 100644
--- a/tensorflow/core/lib/io/buffered_inputstream_test.cc
+++ b/tensorflow/core/lib/io/buffered_inputstream_test.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "tensorflow/core/lib/io/random_inputstream.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
 
 namespace tensorflow {
 namespace io {
@@ -362,6 +363,45 @@ TEST(BufferedInputStream, ReadAll_Text) {
   }
 }
 
+void BM_BufferedReaderSmallReads(const int iters, const int buff_size,
+                                 const int file_size) {
+  testing::StopTiming();
+  Env* env = Env::Default();
+  string fname = testing::TmpDir() + "/buffered_inputstream_test";
+
+  const string file_elem = "0123456789";
+  std::unique_ptr<WritableFile> write_file;
+  TF_ASSERT_OK(env->NewWritableFile(fname, &write_file));
+  for (int i = 0; i < file_size; ++i) {
+    TF_ASSERT_OK(write_file->Append(file_elem));
+  }
+  TF_ASSERT_OK(write_file->Close());
+
+  std::unique_ptr<RandomAccessFile> file;
+  TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
+
+  string result;
+  testing::StartTiming();
+
+  for (int itr = 0; itr < iters; ++itr) {
+    BufferedInputStream in(file.get(), buff_size);
+    for (int64 i = 0; i < 10 * file_size; ++i) {
+      TF_ASSERT_OK(in.ReadNBytes(1, &result))
+          << "i: " << i << " itr: " << itr << " buff_size: " << buff_size
+          << " file size: " << file_size;
+    }
+  }
+}
+BENCHMARK(BM_BufferedReaderSmallReads)
+    ->ArgPair(1, 5)
+    ->ArgPair(1, 1024)
+    ->ArgPair(10, 5)
+    ->ArgPair(10, 1024)
+    ->ArgPair(1024, 1024)
+    ->ArgPair(1024 * 1024, 1024)
+    ->ArgPair(1024 * 1024, 1024 * 1024)
+    ->ArgPair(256 * 1024 * 1024, 1024);
+
 }  // anonymous namespace
 }  // namespace io
 }  // namespace tensorflow