diff --git a/google/cloud/bigtable/row_data.py b/google/cloud/bigtable/row_data.py index 66548be72..0e07ff268 100644 --- a/google/cloud/bigtable/row_data.py +++ b/google/cloud/bigtable/row_data.py @@ -16,6 +16,8 @@ import copy +import time + import six import grpc @@ -342,6 +344,10 @@ def _retry_read_rows_exception(exc): initial=1.0, maximum=15.0, multiplier=2.0, + # NOTE: this is a soft read timeout: this limits for how long we are willing + # to schedule retry attempts to read the next row. This does not set the + # RPC timeout. Please use the separate overal_timeout parameter of read_rows + # to limit the operation duration deadline=60.0, # 60 seconds ) """The default retry strategy to be used on retry-able errors. @@ -389,7 +395,14 @@ class PartialRowsData(object): STATE_CELL_IN_PROGRESS: CELL_IN_PROGRESS, } - def __init__(self, read_method, request, retry=DEFAULT_RETRY_READ_ROWS): + def __init__( + self, + read_method, + request, + retry=DEFAULT_RETRY_READ_ROWS, + attempt_timeout=None, + overall_timeout=None, + ): # Counter for rows returned to the user self._counter = 0 # In-progress row, unset until first response, after commit/reset @@ -406,14 +419,14 @@ def __init__(self, read_method, request, retry=DEFAULT_RETRY_READ_ROWS): self.read_method = read_method self.request = request self.retry = retry + self._attempt_timeout = attempt_timeout + # absolute timestamp when all retry attempts should end + if overall_timeout: + self._overall_deadline = time.time() + overall_timeout + else: + self._overall_deadline = None - # The `timeout` parameter must be somewhat greater than the value - # contained in `self.retry`, in order to avoid race-like condition and - # allow registering the first deadline error before invoking the retry. - # Otherwise there is a risk of entering an infinite loop that resets - # the timeout counter just before it being triggered. The increment - # by 1 second here is customary but should not be much less than that. - self.response_iterator = read_method(request, timeout=self.retry._deadline + 1) + self.response_iterator = self._create_read_stream(request) self.rows = {} self._state = self.STATE_NEW_ROW @@ -451,6 +464,28 @@ class as a generator instead. for row in self: self.rows[row.row_key] = row + @property + def remaining_overall_timeout(self): + """Returns the remaining deadline allotted for the entire stream. + Returns a float of seconds""" + if not self._overall_deadline: + return None + + return self._overall_deadline - time.time() + + def _create_read_stream(self, req): + """Starts a new RPC bounded by the overall deadline and attempt timeout. + + :type req: class:`data_messages_v2_pb2.ReadRowsRequest` + """ + effective_timeout = self.remaining_overall_timeout + if effective_timeout is None: + effective_timeout = self._attempt_timeout + elif self._attempt_timeout is not None: + effective_timeout = min(effective_timeout, self._attempt_timeout) + + return self.read_method(req, timeout=effective_timeout) + def _create_retry_request(self): """Helper for :meth:`__iter__`.""" req_manager = _ReadRowsRequestManager( @@ -465,7 +500,7 @@ def _on_error(self, exc): if self.last_scanned_row_key: retry_request = self._create_retry_request() - self.response_iterator = self.read_method(retry_request) + self.response_iterator = self._create_read_stream(retry_request) def _read_next(self): """Helper for :meth:`__iter__`.""" @@ -476,6 +511,20 @@ def _read_next(self): def _read_next_response(self): """Helper for :meth:`__iter__`.""" + # Calculate the maximum amount of time that retries should be scheduled. + # This will not actually set any deadlines, it will only limit the + # duration of time that we are willing to schedule retries for. + remaining_overall_timeout = self.remaining_overall_timeout + + if remaining_overall_timeout is not None: + # we want make sure that the retry logic doesnt retry after the + # operation deadline is past + if ( + self.retry.deadline is None + or self.retry.deadline > remaining_overall_timeout + ): + self.retry = self.retry.with_deadline(remaining_overall_timeout) + return self.retry(self._read_next, on_error=self._on_error)() def __iter__(self): diff --git a/google/cloud/bigtable/table.py b/google/cloud/bigtable/table.py index 887b74b02..45481f350 100644 --- a/google/cloud/bigtable/table.py +++ b/google/cloud/bigtable/table.py @@ -481,7 +481,7 @@ def get_cluster_states(self): for cluster_id, value_pb in table_pb.cluster_states.items() } - def read_row(self, row_key, filter_=None): + def read_row(self, row_key, filter_=None, overall_timeout=60): """Read a single row from this table. For example: @@ -506,7 +506,11 @@ def read_row(self, row_key, filter_=None): """ row_set = RowSet() row_set.add_row_key(row_key) - result_iter = iter(self.read_rows(filter_=filter_, row_set=row_set)) + result_iter = iter( + self.read_rows( + filter_=filter_, row_set=row_set, overall_timeout=overall_timeout + ) + ) row = next(result_iter, None) if next(result_iter, None) is not None: raise ValueError("More than one row was returned.") @@ -521,6 +525,8 @@ def read_rows( end_inclusive=False, row_set=None, retry=DEFAULT_RETRY_READ_ROWS, + attempt_timeout=None, + overall_timeout=None, ): """Read rows from this table. @@ -565,7 +571,22 @@ def read_rows( default value :attr:`DEFAULT_RETRY_READ_ROWS` can be used and modified with the :meth:`~google.api_core.retry.Retry.with_delay` method or the :meth:`~google.api_core.retry.Retry.with_deadline` - method. + method. This retry object is used to try to fetch the next row: + this means that the deadline specified by this object is reset + after every row read. Furthermore, this deadline is loosely enforced: + it will only prevent additional attempts from be scheduled after the + deadline, it will not limit how long a single attempt to read the + next row will run. Prefer to use overall_timeout below. + + + :type attempt_timeout: float + :param attempt_timeout: (Optional) the attempt timeout to execute a + single RPC. If this attempt fails and there is overall_timeout + left, another attempt will be sent. + + :type overall_timeout: float + :param overall_timeout: (Optional) the overall operation deadline to + to completely read the entire ReadRows stream. :rtype: :class:`.PartialRowsData` :returns: A :class:`.PartialRowsData` a generator for consuming @@ -582,7 +603,13 @@ def read_rows( row_set=row_set, ) data_client = self._instance._client.table_data_client - return PartialRowsData(data_client.transport.read_rows, request_pb, retry) + return PartialRowsData( + data_client.transport.read_rows, + request_pb, + retry, + attempt_timeout=attempt_timeout, + overall_timeout=overall_timeout, + ) def yield_rows(self, **kwargs): """Read rows from this table. @@ -615,6 +642,15 @@ def yield_rows(self, **kwargs): :param row_set: (Optional) The row set containing multiple row keys and row_ranges. + :type attempt_timeout: float + :param attempt_timeout: (Optional) the attempt timeout to execute a + single RPC. If this attempt fails and there is overall_timeout + left, another attempt will be sent. + + :type overall_timeout: float + :param overall_timeout: (Optional) the overall operation deadline to + to completely read the entire ReadRows stream. + :rtype: :class:`.PartialRowData` :returns: A :class:`.PartialRowData` for each row returned """ diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py index 9981380fe..17ab46839 100644 --- a/tests/unit/test_row_data.py +++ b/tests/unit/test_row_data.py @@ -370,20 +370,26 @@ def test_constructor(self): self.assertEqual(partial_rows_data.rows, {}) self.assertEqual(partial_rows_data.retry, DEFAULT_RETRY_READ_ROWS) - def test_constructor_with_retry(self): - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS - + def test_constructor_with_overall_timeout(self): client = _Client() client._data_stub = mock.MagicMock() request = object() - retry = DEFAULT_RETRY_READ_ROWS - partial_rows_data = self._make_one(client._data_stub.ReadRows, request, retry) - partial_rows_data.read_method.assert_called_once_with( - request, timeout=DEFAULT_RETRY_READ_ROWS.deadline + 1 + partial_rows_data = self._make_one( + client._data_stub.ReadRows, request, overall_timeout=11 ) + partial_rows_data.read_method.assert_called_once_with(request, timeout=mock.ANY) + + # the deadline being passed to the first RPC should be slightly less + # than 11. But to avoid flakiness on slow test runners, its padded down + # by 3 secs + self.assertLess(8, partial_rows_data.read_method.call_args.kwargs["timeout"]) + self.assertIs(partial_rows_data.request, request) self.assertEqual(partial_rows_data.rows, {}) - self.assertEqual(partial_rows_data.retry, retry) + # The remaining deadline should be + # But to avoid flakiness on slow test runners, its padded down by 3 secs + self.assertLess(8, partial_rows_data.remaining_overall_timeout) + self.assertLessEqual(partial_rows_data.remaining_overall_timeout, 11) def test___eq__(self): client = _Client() diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 1793a8602..4b42fc3a6 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -11,8 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - - +import time import unittest import mock @@ -798,17 +797,12 @@ def mock_create_row_request(table_name, **kwargs): def test_read_retry_rows(self): from google.cloud.bigtable_v2.gapic import bigtable_client - from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client from google.api_core import retry data_api = bigtable_client.BigtableClient(mock.Mock()) - table_api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock()) credentials = _make_credentials() - client = self._make_client( - project="project-id", credentials=credentials, admin=True - ) + client = self._make_client(project="project-id", credentials=credentials) client._table_data_client = data_api - client._table_admin_client = table_api instance = client.instance(instance_id=self.INSTANCE_ID) table = self._make_one(self.TABLE_ID, instance) @@ -857,6 +851,77 @@ def test_read_retry_rows(self): result = rows[1] self.assertEqual(result.row_key, self.ROW_KEY_2) + def test_read_retry_rows_timeouts(self): + from google.cloud.bigtable_v2.gapic import bigtable_client + + data_api = bigtable_client.BigtableClient(mock.Mock()) + credentials = _make_credentials() + client = self._make_client(project="project-id", credentials=credentials) + client._table_data_client = data_api + instance = client.instance(instance_id=self.INSTANCE_ID) + table = self._make_one(self.TABLE_ID, instance) + + # Patch the stub used by the API method. + client._table_data_client.transport.read_rows = mock.Mock( + side_effect=[_MockReadRowsIterator()] + ) + + # By default there is no timeout + list(table.read_rows()) + self.assertIsNone( + client._table_data_client.transport.read_rows.call_args.kwargs["timeout"] + ) + + # attempt timeout should be passed thru + client._table_data_client.transport.read_rows = mock.Mock( + side_effect=[_MockReadRowsIterator()] + ) + list(table.read_rows(attempt_timeout=1.0)) + self.assertEquals( + 1.0, + client._table_data_client.transport.read_rows.call_args.kwargs["timeout"], + ) + + # overall timeout should be passed thru + client._table_data_client.transport.read_rows = mock.Mock( + side_effect=[_MockReadRowsIterator()] + ) + list(table.read_rows(overall_timeout=10.0)) + # The RPC timeout should be slightly less than 10.0 but to avoid test + # flakiness its padded by a couple of secs. + self.assertLess( + 8.0, + client._table_data_client.transport.read_rows.call_args.kwargs["timeout"], + ) + + # attempt timeout limits overall timeout + client._table_data_client.transport.read_rows = mock.Mock( + side_effect=[_MockReadRowsIterator()] + ) + list(table.read_rows(attempt_timeout=5.0, overall_timeout=10.0)) + self.assertLessEqual( + 5.0, + client._table_data_client.transport.read_rows.call_args.kwargs["timeout"], + ) + + # attempt timeout is truncated by overall timeout + class DelayedFailureIterator(object): + def next(self): + time.sleep(0.75) + raise DeadlineExceeded("delayed error") + + __next__ = next + + client._table_data_client.transport.read_rows = mock.Mock( + side_effect=[DelayedFailureIterator(), _MockReadRowsIterator()] + ) + list(table.read_rows(attempt_timeout=1.0, overall_timeout=1.0)) + + self.assertGreater( + 1.0, + client._table_data_client.transport.read_rows.call_args.kwargs["timeout"], + ) + def test_yield_retry_rows(self): from google.cloud.bigtable_v2.gapic import bigtable_client from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client