Merge pull request #18108 from grpc/enable-load-reporting

Enable Load Reporting

src/python/grpcio/grpc/__init__.py

@@ -14,6 +14,7 @@
 """gRPC's Python API."""
 
 import abc
+import contextlib
 import enum
 import logging
 import sys
@@ -1779,6 +1780,14 @@ def server(thread_pool,
                                  maximum_concurrent_rpcs)
 
 
+@contextlib.contextmanager
+def _create_servicer_context(rpc_event, state, request_deserializer):
+    from grpc import _server  # pylint: disable=cyclic-import
+    context = _server._Context(rpc_event, state, request_deserializer)
+    yield context
+    context._finalize_state()  # pylint: disable=protected-access
+
+
 ###################################  __all__  #################################
 
 __all__ = (

src/python/grpcio/grpc/_server.py

@@ -302,6 +302,9 @@ class _Context(grpc.ServicerContext):
         with self._state.condition:
             self._state.details = _common.encode(details)
 
+    def _finalize_state(self):
+        pass
+
 
 class _RequestIterator(object):
 
@@ -387,16 +390,20 @@ def _unary_request(rpc_event, state, request_deserializer):
 
 
 def _call_behavior(rpc_event, state, behavior, argument, request_deserializer):
-    context = _Context(rpc_event, state, request_deserializer)
+    from grpc import _create_servicer_context
+    with _create_servicer_context(rpc_event, state,
+                                  request_deserializer) as context:
         try:
-        return behavior(argument, context), True
+            response = behavior(argument, context)
+            return response, True
         except Exception as exception:  # pylint: disable=broad-except
             with state.condition:
                 if state.aborted:
                     _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
                            b'RPC Aborted')
                 elif exception not in state.rpc_errors:
-                details = 'Exception calling application: {}'.format(exception)
+                    details = 'Exception calling application: {}'.format(
+                        exception)
                     _LOGGER.exception(details)
                     _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
                            _common.encode(details))