Refactor _plugin_wrapping

src/python/grpcio/grpc/__init__.py

@@ -1156,20 +1156,6 @@ def ssl_channel_credentials(root_certificates=None,
         _cygrpc.channel_credentials_ssl(root_certificates, pair))
 
 
-def _metadata_call_credentials(metadata_plugin, name):
-    from grpc import _plugin_wrapping  # pylint: disable=cyclic-import
-    if name is None:
-        try:
-            effective_name = metadata_plugin.__name__
-        except AttributeError:
-            effective_name = metadata_plugin.__class__.__name__
-    else:
-        effective_name = name
-    return CallCredentials(
-        _plugin_wrapping.call_credentials_metadata_plugin(metadata_plugin,
-                                                          effective_name))
-
-
 def metadata_call_credentials(metadata_plugin, name=None):
     """Construct CallCredentials from an AuthMetadataPlugin.
 
@@ -1180,7 +1166,10 @@ def metadata_call_credentials(metadata_plugin, name=None):
     Returns:
       A CallCredentials.
     """
-    return _metadata_call_credentials(metadata_plugin, name)
+    from grpc import _plugin_wrapping  # pylint: disable=cyclic-import
+    return CallCredentials(
+        _plugin_wrapping.metadata_plugin_call_credentials(metadata_plugin,
+                                                          name))
 
 
 def access_token_call_credentials(access_token):
@@ -1195,8 +1184,10 @@ def access_token_call_credentials(access_token):
       A CallCredentials.
     """
     from grpc import _auth  # pylint: disable=cyclic-import
-    return _metadata_call_credentials(
+    from grpc import _plugin_wrapping  # pylint: disable=cyclic-import
-        _auth.AccessTokenAuthMetadataPlugin(access_token), None)
+    return CallCredentials(
+        _plugin_wrapping.metadata_plugin_call_credentials(
+            _auth.AccessTokenAuthMetadataPlugin(access_token), None))
 
 
 def composite_call_credentials(*call_credentials):

src/python/grpcio/grpc/_plugin_wrapping.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import collections
+import logging
 import threading
 
 import grpc
@@ -20,89 +21,80 @@ from grpc import _common
 from grpc._cython import cygrpc
 
 
-class AuthMetadataContext(
+class _AuthMetadataContext(
         collections.namedtuple('AuthMetadataContext', (
             'service_url', 'method_name',)), grpc.AuthMetadataContext):
     pass
 
 
-class AuthMetadataPluginCallback(grpc.AuthMetadataContext):
+class _CallbackState(object):
 
-    def __init__(self, callback):
+    def __init__(self):
-        self._callback = callback
+        self.lock = threading.Lock()
-
+        self.called = False
-    def __call__(self, metadata, error):
+        self.exception = None
-        self._callback(metadata, error)
 
 
-class _WrappedCygrpcCallback(object):
+class _AuthMetadataPluginCallback(grpc.AuthMetadataPluginCallback):
 
-    def __init__(self, cygrpc_callback):
+    def __init__(self, state, callback):
-        self.is_called = False
+        self._state = state
-        self.error = None
+        self._callback = callback
-        self.is_called_lock = threading.Lock()
-        self.cygrpc_callback = cygrpc_callback
-
-    def _invoke_failure(self, error):
-        # TODO(atash) translate different Exception superclasses into different
-        # status codes.
-        self.cygrpc_callback(_common.EMPTY_METADATA, cygrpc.StatusCode.internal,
-                             _common.encode(str(error)))
-
-    def _invoke_success(self, metadata):
-        try:
-            cygrpc_metadata = _common.to_cygrpc_metadata(metadata)
-        except Exception as exception:  # pylint: disable=broad-except
-            self._invoke_failure(exception)
-            return
-        self.cygrpc_callback(cygrpc_metadata, cygrpc.StatusCode.ok, b'')
 
     def __call__(self, metadata, error):
-        with self.is_called_lock:
+        with self._state.lock:
-            if self.is_called:
+            if self._state.exception is None:
-                raise RuntimeError('callback should only ever be invoked once')
+                if self._state.called:
-            if self.error:
+                    raise RuntimeError(
-                self._invoke_failure(self.error)
+                        'AuthMetadataPluginCallback invoked more than once!')
-                return
+                else:
-            self.is_called = True
+                    self._state.called = True
+            else:
+                raise RuntimeError(
+                    'AuthMetadataPluginCallback raised exception "{}"!'.format(
+                        self._state.exception))
         if error is None:
-            self._invoke_success(metadata)
+            self._callback(
+                _common.to_cygrpc_metadata(metadata), cygrpc.StatusCode.ok, b'')
         else:
-            self._invoke_failure(error)
+            self._callback(_common.EMPTY_METADATA, cygrpc.StatusCode.internal,
-
+                           _common.encode(str(error)))
-    def notify_failure(self, error):
-        with self.is_called_lock:
-            if not self.is_called:
-                self.error = error
 
 
-class _WrappedPlugin(object):
+class _Plugin(object):
 
-    def __init__(self, plugin):
+    def __init__(self, metadata_plugin):
-        self.plugin = plugin
+        self._metadata_plugin = metadata_plugin
 
-    def __call__(self, context, cygrpc_callback):
+    def __call__(self, context, callback):
-        wrapped_cygrpc_callback = _WrappedCygrpcCallback(cygrpc_callback)
+        wrapped_context = _AuthMetadataContext(
-        wrapped_context = AuthMetadataContext(
             _common.decode(context.service_url),
             _common.decode(context.method_name))
+        callback_state = _CallbackState()
+        try:
+            self._metadata_plugin(
+                wrapped_context,
+                _AuthMetadataPluginCallback(callback_state, callback))
+        except Exception as exception:  # pylint: disable=broad-except
+            logging.exception(
+                'AuthMetadataPluginCallback "%s" raised exception!',
+                self._metadata_plugin)
+            with callback_state.lock:
+                callback_state.exception = exception
+                if callback_state.called:
+                    return
+            callback(_common.EMPTY_METADATA, cygrpc.StatusCode.internal,
+                     _common.encode(str(exception)))
+
+
+def metadata_plugin_call_credentials(metadata_plugin, name):
+    if name is None:
         try:
-            self.plugin(wrapped_context,
+            effective_name = metadata_plugin.__name__
-                        AuthMetadataPluginCallback(wrapped_cygrpc_callback))
+        except AttributeError:
-        except Exception as error:
+            effective_name = metadata_plugin.__class__.__name__
-            wrapped_cygrpc_callback.notify_failure(error)
+    else:
-            raise
+        effective_name = name
-
-
-def call_credentials_metadata_plugin(plugin, name):
-    """
-  Args:
-    plugin: A callable accepting a grpc.AuthMetadataContext
-      object and a callback (itself accepting a list of metadata key/value
-      2-tuples and a None-able exception value). The callback must be eventually
-      called, but need not be called in plugin's invocation.
-      plugin's invocation must be non-blocking.
-  """
     return cygrpc.call_credentials_metadata_plugin(
         cygrpc.CredentialsMetadataPlugin(
-            _WrappedPlugin(plugin), _common.encode(name)))
+            _Plugin(metadata_plugin), _common.encode(effective_name)))