Retry test

bigtable/google/cloud/bigtable/retry.py

@@ -0,0 +1,169 @@
+"""Provides function wrappers that implement retrying."""
+import random
+import time
+import six
+
+from google.cloud._helpers import _to_bytes
+from google.cloud.bigtable._generated import (
+ bigtable_pb2 as data_messages_v2_pb2)
+from google.gax import config, errors
+
+_MILLIS_PER_SECOND = 1000
+
+def _has_timeout_settings(backoff_settings):
+ return (backoff_settings.rpc_timeout_multiplier is not None and
+ backoff_settings.max_rpc_timeout_millis is not None and
+ backoff_settings.total_timeout_millis is not None and
+ backoff_settings.initial_rpc_timeout_millis is not None)
+
+class ReadRowsIterator():
+ """Creates an iterator equivalent to a_iter, but that retries on certain
+ exceptions.
+ """
+
+ def __init__(self, client, name, start_key, end_key, filter_, limit,
+ retry_options, **kwargs):
+ self.client = client
+ self.retry_options = retry_options
+ self.name = name
+ self.start_key = start_key
+ self.start_key_closed = True
+ self.end_key = end_key
+ self.filter_ = filter_
+ self.limit = limit
+
+ self.delay_mult = retry_options.backoff_settings.retry_delay_multiplier
+ self.max_delay_millis = retry_options.backoff_settings.max_retry_delay_millis
+ self.has_timeout_settings = _has_timeout_settings(retry_options.backoff_settings)
+
+ if self.has_timeout_settings:
+ self.timeout_mult = retry_options.backoff_settings.rpc_timeout_multiplier
+ self.max_timeout = (retry_options.backoff_settings.max_rpc_timeout_millis / _MILLIS_PER_SECOND)
+ self.total_timeout = (retry_options.backoff_settings.total_timeout_millis / _MILLIS_PER_SECOND)
+ self.set_stream()
+
+ def set_start_key(self, start_key):
+ self.start_key = start_key
+ self.start_key_closed = False
+
+ def set_stream(self):
+ request_pb = _create_row_request(
+ self.name, start_key=self.start_key,
+ start_key_closed=self.start_key_closed, end_key=self.end_key,
+ filter_= self.filter_, limit=self.limit)
+ self.stream = self.client._data_stub.ReadRows(request_pb)
+
+ def next(self, *args, **kwargs):
+ delay = self.retry_options.backoff_settings.initial_retry_delay_millis
+ exc = errors.RetryError('Retry total timeout exceeded before any'
+ 'response was received')
+ if self.has_timeout_settings:
+ timeout = (
+ self.retry_options.backoff_settings.initial_rpc_timeout_millis /
+ _MILLIS_PER_SECOND)
+
+ now = time.time()
+ deadline = now + self.total_timeout
+ else:
+ timeout = None
+ deadline = None
+
+ while deadline is None or now < deadline:
+ try:
+ return six.next(self.stream)
+ except StopIteration as stop:
+ raise stop
+ except Exception as exception: # pylint: disable=broad-except
+ code = config.exc_to_code(exception)
+ if code not in self.retry_options.retry_codes:
+ raise errors.RetryError(
+ 'Exception occurred in retry method that was not'
+ ' classified as transient', exception)
+
+ # pylint: disable=redefined-variable-type
+ exc = errors.RetryError(
+ 'Retry total timeout exceeded with exception', exception)
+
+ # Sleep a random number which will, on average, equal the
+ # expected delay.
+ to_sleep = random.uniform(0, delay * 2)
+ time.sleep(to_sleep / _MILLIS_PER_SECOND)
+ delay = min(delay * self.delay_mult, self.max_delay_millis)
+
+ if self.has_timeout_settings:
+ now = time.time()
+ timeout = min(
+ timeout * self.timeout_mult, self.max_timeout, deadline - now)
+ self.set_stream()
+
+ raise exc
+
+ def __next__(self, *args, **kwargs):
+ return self.next(*args, **kwargs)
+
+ def __iter__(self):
+ return self
+
+def _create_row_request(table_name, row_key=None, start_key=None,
+ start_key_closed=True, end_key=None, filter_=None,
+ limit=None):
+ """Creates a request to read rows in a table.
+
+ :type table_name: str
+ :param table_name: The name of the table to read from.
+
+ :type row_key: bytes
+ :param row_key: (Optional) The key of a specific row to read from.
+
+ :type start_key: bytes
+ :param start_key: (Optional) The beginning of a range of row keys to
+ read from. The range will include ``start_key``. If
+ left empty, will be interpreted as the empty string.
+
+ :type end_key: bytes
+ :param end_key: (Optional) The end of a range of row keys to read from.
+ The range will not include ``end_key``. If left empty,
+ will be interpreted as an infinite string.
+
+ :type filter_: :class:`.RowFilter`
+ :param filter_: (Optional) The filter to apply to the contents of the
+ specified row(s). If unset, reads the entire table.
+
+ :type limit: int
+ :param limit: (Optional) The read will terminate after committing to N
+ rows' worth of results. The default (zero) is to return
+ all results.
+
+ :rtype: :class:`data_messages_v2_pb2.ReadRowsRequest`
+ :returns: The ``ReadRowsRequest`` protobuf corresponding to the inputs.
+ :raises: :class:`ValueError <exceptions.ValueError>` if both
+ ``row_key`` and one of ``start_key`` and ``end_key`` are set
+ """
+ request_kwargs = {'table_name': table_name}
+ if (row_key is not None and
+ (start_key is not None or end_key is not None)):
+ raise ValueError('Row key and row range cannot be '
+ 'set simultaneously')
+ range_kwargs = {}
+ if start_key is not None or end_key is not None:
+ if start_key is not None:
+ if start_key_closed:
+ range_kwargs['start_key_closed'] = _to_bytes(start_key)
+ else:
+ range_kwargs['start_key_open'] = _to_bytes(start_key)
+ if end_key is not None:
+ range_kwargs['end_key_open'] = _to_bytes(end_key)
+ if filter_ is not None:
+ request_kwargs['filter'] = filter_.to_pb()
+ if limit is not None:
+ request_kwargs['rows_limit'] = limit
+
+ message = data_messages_v2_pb2.ReadRowsRequest(**request_kwargs)
+
+ if row_key is not None:
+ message.rows.row_keys.append(_to_bytes(row_key))
+
+ if range_kwargs:
+ message.rows.row_ranges.add(**range_kwargs)
+
+ return message

bigtable/google/cloud/bigtable/row_data.py

@@ -274,6 +274,9 @@ def consume_next(self):
 
  self._validate_chunk(chunk)
 
+ if ("ReadRowsIterator" in self._response_iterator.__class__.__name__):
+ self._response_iterator.set_start_key(chunk.row_key)
+
  if chunk.reset_row:
  row = self._row = None
  cell = self._cell = self._previous_cell = None

bigtable/google/cloud/bigtable/table.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 """User friendly container for Google Cloud Bigtable Table."""
+from __future__ import absolute_import, division
+
+import six
 
-from google.cloud._helpers import _to_bytes
 from google.cloud.bigtable._generated import (
  bigtable_pb2 as data_messages_v2_pb2)
 from google.cloud.bigtable._generated import (
@@ -27,7 +29,29 @@
 from google.cloud.bigtable.row import ConditionalRow
 from google.cloud.bigtable.row import DirectRow
 from google.cloud.bigtable.row_data import PartialRowsData
-
+from google.gax import RetryOptions, BackoffSettings
+from google.cloud.bigtable.retry import ReadRowsIterator, _create_row_request
+from grpc import StatusCode
+
+BACKOFF_SETTINGS = BackoffSettings(
+ initial_retry_delay_millis = 10,
+ retry_delay_multiplier = 2,
+ max_retry_delay_millis = 5000,
+ initial_rpc_timeout_millis = 10,
+ rpc_timeout_multiplier = 2,
+ max_rpc_timeout_millis = 1000,
+ total_timeout_millis = 5000
+)
+
+RETRY_OPTIONS = RetryOptions(
+ retry_codes = [
+ StatusCode.DEADLINE_EXCEEDED,
+ StatusCode.ABORTED,
+ StatusCode.INTERNAL,
+ StatusCode.UNAVAILABLE
+ ],
+ backoff_settings = BACKOFF_SETTINGS
+)
 
 class Table(object):
  """Representation of a Google Cloud Bigtable Table.
@@ -268,13 +292,11 @@ def read_rows(self, start_key=None, end_key=None, limit=None,
  :returns: A :class:`.PartialRowsData` convenience wrapper for consuming
  the streamed results.
  """
- request_pb = _create_row_request(
- self.name, start_key=start_key, end_key=end_key, filter_=filter_,
- limit=limit)
+
  client = self._instance._client
-response_iterator = client._data_stub.ReadRows(request_pb)
-# We expect an iterator of `data_messages_v2_pb2.ReadRowsResponse`
- return PartialRowsData(response_iterator)
+retrying_iterator = ReadRowsIterator(client, self.name, start_key,
+ end_key, filter_, limit, RETRY_OPTIONS)
+ return PartialRowsData(retrying_iterator)
 
  def sample_row_keys(self):
  """Read a sample of row keys in the table.
@@ -312,64 +334,3 @@ def sample_row_keys(self):
  client = self._instance._client
  response_iterator = client._data_stub.SampleRowKeys(request_pb)
  return response_iterator
-
-
-def _create_row_request(table_name, row_key=None, start_key=None, end_key=None,
- filter_=None, limit=None):
- """Creates a request to read rows in a table.
-
- :type table_name: str
- :param table_name: The name of the table to read from.
-
- :type row_key: bytes
- :param row_key: (Optional) The key of a specific row to read from.
-
- :type start_key: bytes
- :param start_key: (Optional) The beginning of a range of row keys to
- read from. The range will include ``start_key``. If
- left empty, will be interpreted as the empty string.
-
- :type end_key: bytes
- :param end_key: (Optional) The end of a range of row keys to read from.
- The range will not include ``end_key``. If left empty,
- will be interpreted as an infinite string.
-
- :type filter_: :class:`.RowFilter`
- :param filter_: (Optional) The filter to apply to the contents of the
- specified row(s). If unset, reads the entire table.
-
- :type limit: int
- :param limit: (Optional) The read will terminate after committing to N
- rows' worth of results. The default (zero) is to return
- all results.
-
- :rtype: :class:`data_messages_v2_pb2.ReadRowsRequest`
- :returns: The ``ReadRowsRequest`` protobuf corresponding to the inputs.
- :raises: :class:`ValueError <exceptions.ValueError>` if both
- ``row_key`` and one of ``start_key`` and ``end_key`` are set
- """
- request_kwargs = {'table_name': table_name}
- if (row_key is not None and
- (start_key is not None or end_key is not None)):
- raise ValueError('Row key and row range cannot be '
- 'set simultaneously')
- range_kwargs = {}
- if start_key is not None or end_key is not None:
- if start_key is not None:
- range_kwargs['start_key_closed'] = _to_bytes(start_key)
- if end_key is not None:
- range_kwargs['end_key_open'] = _to_bytes(end_key)
- if filter_ is not None:
- request_kwargs['filter'] = filter_.to_pb()
- if limit is not None:
- request_kwargs['rows_limit'] = limit
-
- message = data_messages_v2_pb2.ReadRowsRequest(**request_kwargs)
-
- if row_key is not None:
- message.rows.row_keys.append(_to_bytes(row_key))
-
- if range_kwargs:
- message.rows.row_ranges.add(**range_kwargs)
-
- return message