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