fix(bigtable): ReadRows retries from the last_scanned_row_key (googleapis#8423)

Author: dbolduc

google/cloud/bigtable/async_row_reader.h

@@ -374,6 +374,9 @@ class AsyncRowReader : public std::enable_shared_from_this<
         return parser_status;
       }
     }
+    if (!response.last_scanned_row_key().empty()) {
+      last_read_row_key_ = std::move(*response.mutable_last_scanned_row_key());
+    }
     return Status();
   }
 

google/cloud/bigtable/async_row_reader_test.cc

@@ -662,6 +662,82 @@ TEST_F(TableAsyncReadRowsTest, TransientErrorIsRetried) {
   ASSERT_EQ(0U, cq_impl_->size());
 }
 
+/// @test Verify that the last scanned row is respected.
+TEST_F(TableAsyncReadRowsTest, LastScannedRowKeyIsRespected) {
+  auto& stream2 = AddReader([](btproto::ReadRowsRequest const& req) {
+    // The server has told that "r2" has been scanned. Our second request
+    // should use the range ("r2", ""]. This is what is under test.
+    EXPECT_TRUE(req.has_rows());
+    auto const& rows = req.rows();
+    EXPECT_EQ(1, rows.row_ranges_size());
+    auto const& range = rows.row_ranges(0);
+    EXPECT_EQ("r2", range.start_key_open());
+  });
+  auto& stream1 = AddReader([](btproto::ReadRowsRequest const&) {});
+
+  EXPECT_CALL(stream1, Read)
+      .WillOnce([](btproto::ReadRowsResponse* r, void*) {
+        *r = bigtable::testing::ReadRowsResponseFromString(
+            R"(
+                chunks {
+                  row_key: "r1"
+                  family_name { value: "fam" }
+                  qualifier { value: "col" }
+                  timestamp_micros: 42000
+                  value: "value"
+                  commit_row: true
+                })");
+      })
+      .WillOnce([](btproto::ReadRowsResponse* r, void*) {
+        std::cout << "Second Read for stream1 " << std::endl;
+        btproto::ReadRowsResponse resp;
+        resp.set_last_scanned_row_key("r2");
+        *r = resp;
+      })
+      .RetiresOnSaturation();
+  EXPECT_CALL(stream1, Finish).WillOnce([](grpc::Status* status, void*) {
+    *status = grpc::Status(grpc::StatusCode::UNAVAILABLE, "retry");
+  });
+
+  EXPECT_CALL(stream2, Finish).WillOnce([](grpc::Status* status, void*) {
+    *status = grpc::Status::OK;
+  });
+
+  ExpectRows({"r1", "r2", "r3"});
+  promises_from_user_cb_[0].set_value(true);
+  ReadRows();
+
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(true);  // Finish Start()
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(true);  // Return "r1"
+
+  row_futures_[0].get();
+
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(true);  // Return last_scanned_row_key = "r2"
+
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(false);  // Finish stream1 with failure
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(true);  // Finish Finish()
+
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(true);  // Finish timer
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(true);  // Finish Start()
+
+  ASSERT_EQ(1U, cq_impl_->size());
+  cq_impl_->SimulateCompletion(false);  // Finish stream
+  ASSERT_EQ(1U, cq_impl_->size());
+  EXPECT_TRUE(Unsatisfied(stream_status_future_));
+  cq_impl_->SimulateCompletion(true);  // Finish Finish()
+
+  auto stream_status = stream_status_future_.get();
+  ASSERT_STATUS_OK(stream_status);
+  ASSERT_EQ(0U, cq_impl_->size());
+}
+
 /// @test Verify proper handling of bogus responses from the service.
 TEST_F(TableAsyncReadRowsTest, ParserFailure) {
   auto& stream = AddReader([](btproto::ReadRowsRequest const&) {});

google/cloud/bigtable/row_reader.cc

@@ -140,6 +140,9 @@ bool RowReader::NextChunk() {
       response_ = {};
       return false;
     }
+    if (!response_.last_scanned_row_key().empty()) {
+      last_read_row_key_ = std::move(*response_.mutable_last_scanned_row_key());
+    }
   }
   return true;
 }

google/cloud/bigtable/row_reader_test.cc

@@ -394,6 +394,61 @@ TEST_F(RowReaderTest, FailedStreamRetriesSkipAlreadyReadRows) {
   EXPECT_EQ(++it, reader.end());
 }
 
+TEST_F(RowReaderTest, FailedStreamRetriesSkipToLastScannedRow) {
+  auto* stream = new MockReadRowsReader("google.bigtable.v2.Bigtable.ReadRows");
+  auto parser = absl::make_unique<ReadRowsParserMock>();
+  parser->SetRows({"r1"});
+  google::bigtable::v2::ReadRowsResponse response;
+  response.set_last_scanned_row_key("r2");
+  {
+    ::testing::InSequence s;
+    // We start our call with 3 rows in the set: "r1", "r2", "r3".
+    EXPECT_CALL(*client_, ReadRows(_, RequestWithRowKeysCount(3)))
+        .WillOnce(stream->MakeMockReturner());
+
+    // The mock `parser` will return "r1". Next, simulate the server returning
+    // an empty chunk with `last_scanned_row_key` set to "r2".
+    EXPECT_CALL(*stream, Read)
+        .WillOnce(DoAll(SetArgPointee<0>(response), Return(true)));
+
+    // The stream fails with a retry-able error
+    EXPECT_CALL(*stream, Read).WillOnce(Return(false));
+    EXPECT_CALL(*stream, Finish())
+        .WillOnce(Return(grpc::Status(grpc::StatusCode::INTERNAL, "retry")));
+
+    EXPECT_CALL(*retry_policy_, OnFailure(An<grpc::Status const&>()))
+        .WillOnce(Return(true));
+    EXPECT_CALL(*backoff_policy_, OnCompletion(An<grpc::Status const&>()))
+        .WillOnce(Return(std::chrono::milliseconds(0)));
+
+    auto* stream_retry =
+        new MockReadRowsReader("google.bigtable.v2.Bigtable.ReadRows");
+
+    // We retry the remaining rows. We have "r1" returned, but the service has
+    // also told us that "r2" was scanned. This means there is only one row
+    // remaining to read: "r3".
+    EXPECT_CALL(*client_, ReadRows(_, RequestWithRowKeysCount(1)))
+        .WillOnce(stream_retry->MakeMockReturner());
+
+    // End the stream to clean up the test
+    EXPECT_CALL(*stream_retry, Read).WillOnce(Return(false));
+    EXPECT_CALL(*stream_retry, Finish()).WillOnce(Return(grpc::Status::OK));
+  }
+
+  parser_factory_->AddParser(std::move(parser));
+  bigtable::RowReader reader(
+      client_, "", bigtable::RowSet("r1", "r2", "r3"),
+      bigtable::RowReader::NO_ROWS_LIMIT, bigtable::Filter::PassAllFilter(),
+      std::move(retry_policy_), std::move(backoff_policy_),
+      metadata_update_policy_, std::move(parser_factory_));
+
+  auto it = reader.begin();
+  EXPECT_NE(it, reader.end());
+  ASSERT_STATUS_OK(*it);
+  EXPECT_EQ((*it)->row_key(), "r1");
+  EXPECT_EQ(++it, reader.end());
+}
+
 TEST_F(RowReaderTest, FailedParseIsRetried) {
   // wrapped in unique_ptr by ReadRows
   auto* stream = new MockReadRowsReader("google.bigtable.v2.Bigtable.ReadRows");