chore: increase async rpc performance (#1755)

Author: daniel-sanche

gapic/templates/%namespace/%name_%version/%sub/services/%service/_client_macros.j2

@@ -87,8 +87,8 @@
         {% if not method.client_streaming %}
         # Create or coerce a protobuf request object.
         {% if method.flattened_fields %}
-        # Quick check: If we got a request object, we should *not* have
-        # gotten any keyword arguments that map to the request.
+        # - Quick check: If we got a request object, we should *not* have
+        #   gotten any keyword arguments that map to the request.
         has_flattened_params = any([{{ method.flattened_fields.values()|join(", ", attribute="name") }}])
         if request is not None and has_flattened_params:
             raise ValueError('If the `request` argument is set, then none of '
@@ -97,17 +97,15 @@
         {% endif %}
         {% if method.input.ident.package != method.ident.package %}{# request lives in a different package, so there is no proto wrapper #}
         if isinstance(request, dict):
-            # The request isn't a proto-plus wrapped type,
-            # so it must be constructed via keyword expansion.
+            # - The request isn't a proto-plus wrapped type,
+            #   so it must be constructed via keyword expansion.
             request = {{ method.input.ident }}(**request)
         elif not request:
             # Null request, just make one.
             request = {{ method.input.ident }}()
         {% else %}
-        # Minor optimization to avoid making a copy if the user passes
-        # in a {{ method.input.ident }}.
-        # There's no risk of modifying the input as we've already verified
-        # there are no flattened fields.
+        # - Use the request object if provided (there's no risk of modifying the input as
+        #   there are no flattened fields), or create one.
         if not isinstance(request, {{ method.input.ident }}):
             request = {{ method.input.ident }}(request)
         {% endif %}{# different request package #}

gapic/templates/%namespace/%name_%version/%sub/services/%service/async_client.py.j2

@@ -314,23 +314,26 @@ class {{ service.async_client_name }}:
         {% if not method.client_streaming %}
         # Create or coerce a protobuf request object.
         {% if method.flattened_fields %}
-        # Quick check: If we got a request object, we should *not* have
-        # gotten any keyword arguments that map to the request.
+        # - Quick check: If we got a request object, we should *not* have
+        #   gotten any keyword arguments that map to the request.
         has_flattened_params = any([{{ method.flattened_fields.values()|join(", ", attribute="name") }}])
         if request is not None and has_flattened_params:
             raise ValueError("If the `request` argument is set, then none of "
                              "the individual field arguments should be set.")
 
         {% endif %}
         {% if method.input.ident.package != method.ident.package %} {# request lives in a different package, so there is no proto wrapper #}
-        # The request isn't a proto-plus wrapped type,
-        # so it must be constructed via keyword expansion.
+        # - The request isn't a proto-plus wrapped type,
+        #   so it must be constructed via keyword expansion.
         if isinstance(request, dict):
             request = {{ method.input.ident }}(**request)
         elif not request:
             request = {{ method.input.ident }}({% if method.flattened_fields %}{% for f in method.flattened_fields.values() %}{{ f.name }}={{ f.name }}{% if not loop.last %}, {% endif %}{% endfor %}{% endif %})
         {% else %}
-        request = {{ method.input.ident }}(request)
+        # - Use the request object if provided (there's no risk of modifying the input as
+        #   there are no flattened fields), or create one.
+        if not isinstance(request, {{ method.input.ident }}):
+            request = {{ method.input.ident }}(request)
         {% endif %} {# different request package #}
 
         {# Vanilla python protobuf wrapper types cannot _set_ repeated fields #}
@@ -357,26 +360,9 @@ class {{ service.async_client_name }}:
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.{{ method.transport_safe_name|snake_case }},
-            {% if method.retry %}
-            default_retry=retries.AsyncRetry(
-                {% if method.retry.initial_backoff %}initial={{ method.retry.initial_backoff }},{% endif %}
-                {% if method.retry.max_backoff %}maximum={{ method.retry.max_backoff }},{% endif %}
-                {% if method.retry.backoff_multiplier %}multiplier={{ method.retry.backoff_multiplier }},{% endif %}
-                predicate=retries.if_exception_type(
-                    {% for ex in method.retry.retryable_exceptions|sort(attribute="__name__") %}
-                    core_exceptions.{{ ex.__name__ }},
-                    {% endfor %}
-                ),
-                deadline={{ method.timeout }},
-            ),
-            {% endif %}
-            default_timeout={{ method.timeout }},
-            client_info=DEFAULT_CLIENT_INFO,
-        )
-        {% if method.field_headers %}
+        rpc = self._client._transport._wrapped_methods[self._client._transport.{{ method.transport_safe_name|snake_case }}]
 
+        {% if method.field_headers %}
         # Certain fields should be provided within the metadata header;
         # add these here.
         metadata = tuple(metadata) + (

gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/base.py.j2

@@ -149,9 +149,15 @@ class {{ service.name }}Transport(abc.ABC):
                 self.{{ method.transport_safe_name|snake_case }},
                 {% if method.retry %}
                 default_retry=retries.Retry(
-                    {% if method.retry.initial_backoff %}initial={{ method.retry.initial_backoff }},{% endif %}
-                    {% if method.retry.max_backoff %}maximum={{ method.retry.max_backoff }},{% endif %}
-                    {% if method.retry.backoff_multiplier %}multiplier={{ method.retry.backoff_multiplier }},{% endif %}
+                    {% if method.retry.initial_backoff %}
+                    initial={{ method.retry.initial_backoff }},
+                    {% endif %}
+                    {% if method.retry.max_backoff %}
+                    maximum={{ method.retry.max_backoff }},
+                    {% endif %}
+                    {% if method.retry.backoff_multiplier %}
+                    multiplier={{ method.retry.backoff_multiplier }},
+                    {% endif %}
                     predicate=retries.if_exception_type(
                         {% for ex in method.retry.retryable_exceptions|sort(attribute='__name__') %}
                         core_exceptions.{{ ex.__name__ }},

gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/grpc_asyncio.py.j2

@@ -7,6 +7,8 @@ from typing import Awaitable, Callable, Dict, Optional, Sequence, Tuple, Union
 
 from google.api_core import gapic_v1
 from google.api_core import grpc_helpers_async
+from google.api_core import exceptions as core_exceptions
+from google.api_core import retry_async as retries
 {% if service.has_lro %}
 from google.api_core import operations_v1
 {% endif %}
@@ -382,6 +384,37 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
         return self._stubs["test_iam_permissions"]
     {% endif %}
 
+    def _prep_wrapped_messages(self, client_info):
+        """ Precompute the wrapped methods, overriding the base class method to use async wrappers."""
+        self._wrapped_methods = {
+            {% for method in service.methods.values() %}
+            self.{{ method.transport_safe_name|snake_case }}: gapic_v1.method_async.wrap_method(
+                self.{{ method.transport_safe_name|snake_case }},
+                {% if method.retry %}
+                default_retry=retries.AsyncRetry(
+                    {% if method.retry.initial_backoff %}
+                    initial={{ method.retry.initial_backoff }},
+                    {% endif %}
+                    {% if method.retry.max_backoff %}
+                    maximum={{ method.retry.max_backoff }},
+                    {% endif %}
+                    {% if method.retry.backoff_multiplier %}
+                    multiplier={{ method.retry.backoff_multiplier }},
+                    {% endif %}
+                    predicate=retries.if_exception_type(
+                        {% for ex in method.retry.retryable_exceptions|sort(attribute='__name__') %}
+                        core_exceptions.{{ ex.__name__ }},
+                        {% endfor %}
+                    ),
+                    deadline={{ method.timeout }},
+                ),
+                {% endif %}
+                default_timeout={{ method.timeout }},
+                client_info=client_info,
+            ),
+            {% endfor %} {# service.methods.values() #}
+        }
+
     def close(self):
         return self.grpc_channel.close()
 

gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2

@@ -196,6 +196,54 @@ def test_{{ method_name }}_non_empty_request_with_auto_populated_field():
         )
 {% endif %}
 
+def test_{{ method_name }}_use_cached_wrapped_rpc():
+    # Clients should use _prep_wrapped_messages to create cached wrapped rpcs,
+    # instead of constructing them on each call
+    with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn:
+        client = {{ service.client_name }}(
+            credentials=ga_credentials.AnonymousCredentials(),
+            transport="grpc",
+        )
+
+        # Should wrap all calls on client creation
+        assert wrapper_fn.call_count > 0
+        wrapper_fn.reset_mock()
+
+        # Ensure method has been cached
+        assert client._transport.{{method.transport_safe_name|snake_case}} in client._transport._wrapped_methods
+
+        # Replace cached wrapped function with mock
+        mock_rpc = mock.Mock()
+        client._transport._wrapped_methods[client._transport.{{method.transport_safe_name|snake_case}}] = mock_rpc
+
+        {% if method.client_streaming %}
+        request = [{}]
+        client.{{ method.safe_name|snake_case }}(request)
+        {% else %}
+        request = {}
+        client.{{ method_name }}(request)
+        {% endif %}
+
+        # Establish that the underlying gRPC stub method was called.
+        assert mock_rpc.call_count == 1
+
+        {% if method.lro or method.extended_lro %}
+        # Operation methods build a cached wrapper on first rpc call
+        # subsequent calls should use the cached wrapper
+        wrapper_fn.reset_mock()
+        {% endif %}
+
+        {% if method.client_streaming %}
+        client.{{ method.safe_name|snake_case }}(request)
+        {% else %}
+        client.{{ method_name }}(request)
+        {% endif %}
+
+
+        # Establish that a new wrapper was not created for this call
+        assert wrapper_fn.call_count == 0
+        assert mock_rpc.call_count == 2
+
 {% if not full_extended_lro %}
 {% if not method.client_streaming %}
 @pytest.mark.asyncio
@@ -253,6 +301,58 @@ async def test_{{ method_name }}_empty_call_async():
         assert args[0] == {{ method.input.ident }}()
 {% endif %}
 
+@pytest.mark.asyncio
+async def test_{{ method_name }}_async_use_cached_wrapped_rpc(transport: str = "grpc_asyncio"):
+    # Clients should use _prep_wrapped_messages to create cached wrapped rpcs,
+    # instead of constructing them on each call
+    with mock.patch("google.api_core.gapic_v1.method_async.wrap_method") as wrapper_fn:
+        client = {{ service.async_client_name }}(
+            credentials=ga_credentials.AnonymousCredentials(),
+            transport=transport,
+        )
+
+        # Should wrap all calls on client creation
+        assert wrapper_fn.call_count > 0
+        wrapper_fn.reset_mock()
+
+        # Ensure method has been cached
+        assert client._client._transport.{{method.transport_safe_name|snake_case}} in client._client._transport._wrapped_methods
+
+        # Replace cached wrapped function with mock
+        class AwaitableMock(mock.AsyncMock):
+            def __await__(self):
+                self.await_count += 1
+                return iter([])
+        mock_object = AwaitableMock()
+        client._client._transport._wrapped_methods[client._client._transport.{{method.transport_safe_name|snake_case}}] = mock_object
+
+        {% if method.client_streaming %}
+        request = [{}]
+        await client.{{ method.name|snake_case }}(request)
+        {% else %}
+        request = {}
+        await client.{{ method_name }}(request)
+        {% endif %}
+
+        # Establish that the underlying gRPC stub method was called.
+        assert mock_object.call_count == 1
+
+        {% if method.lro or method.extended_lro %}
+        # Operation methods build a cached wrapper on first rpc call
+        # subsequent calls should use the cached wrapper
+        wrapper_fn.reset_mock()
+        {% endif %}
+
+        {% if method.client_streaming %}
+        await client.{{ method.name|snake_case }}(request)
+        {% else %}
+        await client.{{ method_name }}(request)
+        {% endif %}
+
+        # Establish that a new wrapper was not created for this call
+        assert wrapper_fn.call_count == 0
+        assert mock_object.call_count == 2
+
 @pytest.mark.asyncio
 async def test_{{ method_name }}_async(transport: str = 'grpc_asyncio', request_type={{ method.input.ident }}):
     {% with auto_populated_field_sample_value = "explicit value for autopopulate-able field" %}
@@ -1220,6 +1320,53 @@ def test_{{ method_name }}_rest(request_type):
     {% endfor %}
     {% endif %}
 
+def test_{{ method_name }}_rest_use_cached_wrapped_rpc():
+    # Clients should use _prep_wrapped_messages to create cached wrapped rpcs,
+    # instead of constructing them on each call
+    with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn:
+        client = {{ service.client_name }}(
+            credentials=ga_credentials.AnonymousCredentials(),
+            transport="rest",
+        )
+
+        # Should wrap all calls on client creation
+        assert wrapper_fn.call_count > 0
+        wrapper_fn.reset_mock()
+
+        # Ensure method has been cached
+        assert client._transport.{{method.transport_safe_name|snake_case}} in client._transport._wrapped_methods
+
+        # Replace cached wrapped function with mock
+        mock_rpc = mock.Mock()
+        client._transport._wrapped_methods[client._transport.{{method.transport_safe_name|snake_case}}] = mock_rpc
+
+        {% if method.client_streaming %}
+        request = [{}]
+        client.{{ method.safe_name|snake_case }}(request)
+        {% else %}
+        request = {}
+        client.{{ method_name }}(request)
+        {% endif %}
+
+        # Establish that the underlying gRPC stub method was called.
+        assert mock_rpc.call_count == 1
+
+        {% if method.lro or method.extended_lro %}
+        # Operation methods build a cached wrapper on first rpc call
+        # subsequent calls should use the cached wrapper
+        wrapper_fn.reset_mock()
+        {% endif %}
+
+        {% if method.client_streaming %}
+        client.{{ method.safe_name|snake_case }}(request)
+        {% else %}
+        client.{{ method_name }}(request)
+        {% endif %}
+
+        # Establish that a new wrapper was not created for this call
+        assert wrapper_fn.call_count == 0
+        assert mock_rpc.call_count == 2
+
 
     {% if method.input.required_fields %}
 def test_{{ method_name }}_rest_required_fields(request_type={{ method.input.ident }}):