fix: grpc channel refresh (#1087)

* added failing test * prevent _start_bg_task from running * let bg task run * invalidate transport stubs after channel refresh * added sync implementation

Author: daniel-sanche

google/cloud/bigtable/data/_async/client.py

@@ -164,8 +164,8 @@ def __init__(
  if "pool_size" in kwargs:
  warnings.warn("pool_size no longer supported")
  # set up client info headers for veneer library
- client_info = DEFAULT_CLIENT_INFO
- client_info.client_library_version = self._client_version()
+ self.client_info = DEFAULT_CLIENT_INFO
+ self.client_info.client_library_version = self._client_version()
  # parse client options
  if type(client_options) is dict:
  client_options = client_options_lib.from_dict(client_options)
@@ -196,7 +196,7 @@ def __init__(
  self._gapic_client = CrossSync.GapicClient(
  credentials=credentials,
  client_options=client_options,
- client_info=client_info,
+ client_info=self.client_info,
  transport=lambda *args, **kwargs: TransportType(
  *args, **kwargs, channel=custom_channel
  ),
@@ -371,6 +371,9 @@ async def _manage_channel(
  await self._ping_and_warm_instances(channel=new_channel)
  # cycle channel out of use, with long grace window before closure
  self.transport._grpc_channel = new_channel
+ # invalidate caches
+ self.transport._stubs = {}
+ self.transport._prep_wrapped_messages(self.client_info)
  # give old_channel a chance to complete existing rpcs
  if CrossSync.is_async:
  await old_channel.close(grace_period)

google/cloud/bigtable/data/_sync_autogen/client.py

@@ -114,8 +114,8 @@ def __init__(
  """
  if "pool_size" in kwargs:
  warnings.warn("pool_size no longer supported")
- client_info = DEFAULT_CLIENT_INFO
- client_info.client_library_version = self._client_version()
+ self.client_info = DEFAULT_CLIENT_INFO
+ self.client_info.client_library_version = self._client_version()
  if type(client_options) is dict:
  client_options = client_options_lib.from_dict(client_options)
  client_options = cast(
@@ -143,7 +143,7 @@ def __init__(
  self._gapic_client = CrossSync._Sync_Impl.GapicClient(
  credentials=credentials,
  client_options=client_options,
- client_info=client_info,
+ client_info=self.client_info,
  transport=lambda *args, **kwargs: TransportType(
  *args, **kwargs, channel=custom_channel
  ),
@@ -284,6 +284,8 @@ def _manage_channel(
  new_channel = self.transport.create_channel()
  self._ping_and_warm_instances(channel=new_channel)
  self.transport._grpc_channel = new_channel
+ self.transport._stubs = {}
+ self.transport._prep_wrapped_messages(self.client_info)
  if grace_period:
  self._is_closed.wait(grace_period)
  old_channel.close()

tests/system/data/test_system_async.py

@@ -207,6 +207,38 @@ async def test_ping_and_warm(self, client, table):
  assert len(results) == 1
  assert results[0] is None
 
+ @CrossSync.pytest
+ async def test_channel_refresh(self, table_id, instance_id, temp_rows):
+ """
+ change grpc channel to refresh after 1 second. Schedule a read_rows call after refresh,
+ to ensure new channel works
+ """
+ await temp_rows.add_row(b"row_key_1")
+ await temp_rows.add_row(b"row_key_2")
+ project = os.getenv("GOOGLE_CLOUD_PROJECT") or None
+ client = CrossSync.DataClient(project=project)
+ # start custom refresh task
+ try:
+ client._channel_refresh_task = CrossSync.create_task(
+ client._manage_channel,
+ refresh_interval_min=1,
+ refresh_interval_max=1,
+ sync_executor=client._executor,
+ )
+ # let task run
+ await CrossSync.yield_to_event_loop()
+ async with client.get_table(instance_id, table_id) as table:
+ rows = await table.read_rows({})
+ first_channel = client.transport.grpc_channel
+ assert len(rows) == 2
+ await CrossSync.sleep(2)
+ rows_after_refresh = await table.read_rows({})
+ assert len(rows_after_refresh) == 2
+ assert client.transport.grpc_channel is not first_channel
+ print(table)
+ finally:
+ await client.close()
+
  @CrossSync.pytest
  @pytest.mark.usefixtures("table")
  @CrossSync.Retry(

tests/system/data/test_system_autogen.py

@@ -165,6 +165,33 @@ def test_ping_and_warm(self, client, table):
  assert len(results) == 1
  assert results[0] is None
 
+ def test_channel_refresh(self, table_id, instance_id, temp_rows):
+ """change grpc channel to refresh after 1 second. Schedule a read_rows call after refresh,
+ to ensure new channel works"""
+ temp_rows.add_row(b"row_key_1")
+ temp_rows.add_row(b"row_key_2")
+ project = os.getenv("GOOGLE_CLOUD_PROJECT") or None
+ client = CrossSync._Sync_Impl.DataClient(project=project)
+ try:
+ client._channel_refresh_task = CrossSync._Sync_Impl.create_task(
+ client._manage_channel,
+ refresh_interval_min=1,
+ refresh_interval_max=1,
+ sync_executor=client._executor,
+ )
+ CrossSync._Sync_Impl.yield_to_event_loop()
+ with client.get_table(instance_id, table_id) as table:
+ rows = table.read_rows({})
+ first_channel = client.transport.grpc_channel
+ assert len(rows) == 2
+ CrossSync._Sync_Impl.sleep(2)
+ rows_after_refresh = table.read_rows({})
+ assert len(rows_after_refresh) == 2
+ assert client.transport.grpc_channel is not first_channel
+ print(table)
+ finally:
+ client.close()
+
  @pytest.mark.usefixtures("table")
  @CrossSync._Sync_Impl.Retry(
  predicate=retry.if_exception_type(ClientError), initial=1, maximum=5