Cache the out-of-range status and avoid repeating reads in BufferedInputStream.
PiperOrigin-RevId: 302541219 Change-Id: I0f624d6605ddb9aa21a24b809fff1d5de7c6d76e
This commit is contained in:
parent
a58b1026bd
commit
3e2b031211
@ -49,8 +49,7 @@ Status BufferedInputStream::FillBuffer() {
|
||||
Status s = input_stream_->ReadNBytes(size_, &buf_);
|
||||
pos_ = 0;
|
||||
limit_ = buf_.size();
|
||||
if (buf_.empty()) {
|
||||
DCHECK(!s.ok());
|
||||
if (!s.ok()) {
|
||||
file_status_ = s;
|
||||
}
|
||||
return s;
|
||||
@ -93,7 +92,7 @@ Status BufferedInputStream::ReadNBytes(int64 bytes_to_read, tstring* result) {
|
||||
bytes_to_read);
|
||||
}
|
||||
result->clear();
|
||||
if (!file_status_.ok() && bytes_to_read > 0) {
|
||||
if (pos_ == limit_ && !file_status_.ok() && bytes_to_read > 0) {
|
||||
return file_status_;
|
||||
}
|
||||
result->reserve(bytes_to_read);
|
||||
|
||||
@ -30,6 +30,37 @@ static std::vector<int> BufferSizes() {
|
||||
12, 13, 14, 15, 16, 17, 18, 19, 20, 65536};
|
||||
}
|
||||
|
||||
// This class will only return OutOfRange error once to make sure that
|
||||
// BufferedInputStream is able to cache the error.
|
||||
class ReadOnceInputStream : public InputStreamInterface {
|
||||
public:
|
||||
ReadOnceInputStream() : start_(true) {}
|
||||
|
||||
virtual Status ReadNBytes(int64 bytes_to_read, tstring* result) {
|
||||
if (bytes_to_read < 11) {
|
||||
return errors::InvalidArgument("Not reading all bytes: ", bytes_to_read);
|
||||
}
|
||||
if (start_) {
|
||||
*result = "0123456789";
|
||||
start_ = false;
|
||||
return errors::OutOfRange("Out of range.");
|
||||
}
|
||||
return errors::InvalidArgument(
|
||||
"Redudant call to ReadNBytes after an OutOfRange error.");
|
||||
}
|
||||
|
||||
int64 Tell() const override { return start_ ? 0 : 10; }
|
||||
|
||||
// Resets the stream to the beginning.
|
||||
Status Reset() override {
|
||||
start_ = true;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
private:
|
||||
bool start_;
|
||||
};
|
||||
|
||||
TEST(BufferedInputStream, ReadLine_Empty) {
|
||||
Env* env = Env::Default();
|
||||
string fname;
|
||||
@ -196,6 +227,32 @@ TEST(BufferedInputStream, ReadNBytes) {
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BufferedInputStream, OutOfRangeCache) {
|
||||
for (auto buf_size : BufferSizes()) {
|
||||
if (buf_size < 11) {
|
||||
continue;
|
||||
}
|
||||
ReadOnceInputStream input_stream;
|
||||
tstring read;
|
||||
BufferedInputStream in(&input_stream, buf_size);
|
||||
EXPECT_EQ(0, in.Tell());
|
||||
TF_ASSERT_OK(in.ReadNBytes(3, &read));
|
||||
EXPECT_EQ(read, "012");
|
||||
EXPECT_EQ(3, in.Tell());
|
||||
TF_ASSERT_OK((in.ReadNBytes(7, &read)));
|
||||
EXPECT_EQ(read, "3456789");
|
||||
EXPECT_EQ(10, in.Tell());
|
||||
Status s = in.ReadNBytes(5, &read);
|
||||
// Make sure the read is failing with OUT_OF_RANGE error. If it is failing
|
||||
// with other errors, it is not caching the OUT_OF_RANGE properly.
|
||||
EXPECT_EQ(error::OUT_OF_RANGE, s.code()) << s;
|
||||
EXPECT_EQ(read, "");
|
||||
// Empty read shouldn't cause an error even at the end of the file.
|
||||
TF_ASSERT_OK(in.ReadNBytes(0, &read));
|
||||
EXPECT_EQ(read, "");
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BufferedInputStream, SkipNBytes) {
|
||||
Env* env = Env::Default();
|
||||
string fname;
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user