|
|
@@ -13,37 +13,130 @@
|
|
|
# limitations under the License.
|
|
|
"""Tests server and client side compression."""
|
|
|
|
|
|
+from __future__ import absolute_import
|
|
|
+from __future__ import division
|
|
|
+from __future__ import print_function
|
|
|
+
|
|
|
import unittest
|
|
|
|
|
|
+import contextlib
|
|
|
+from concurrent import futures
|
|
|
+import functools
|
|
|
+import itertools
|
|
|
import logging
|
|
|
+import os
|
|
|
+
|
|
|
import grpc
|
|
|
from grpc import _grpcio_metadata
|
|
|
|
|
|
from tests.unit import test_common
|
|
|
from tests.unit.framework.common import test_constants
|
|
|
+from tests.unit import _tcp_proxy
|
|
|
|
|
|
_UNARY_UNARY = '/test/UnaryUnary'
|
|
|
+_UNARY_STREAM = '/test/UnaryStream'
|
|
|
+_STREAM_UNARY = '/test/StreamUnary'
|
|
|
_STREAM_STREAM = '/test/StreamStream'
|
|
|
|
|
|
+# Cut down on test time.
|
|
|
+_STREAM_LENGTH = test_constants.STREAM_LENGTH // 8
|
|
|
+
|
|
|
+_HOST = 'localhost'
|
|
|
+
|
|
|
+_REQUEST = b'\x00' * 100
|
|
|
+_COMPRESSION_RATIO_THRESHOLD = 0.1
|
|
|
+_COMPRESSION_METHODS = (
|
|
|
+ None,
|
|
|
+ # Disabled for test tractability.
|
|
|
+ # grpc.Compression.NoCompression,
|
|
|
+ grpc.Compression.Deflate,
|
|
|
+ grpc.Compression.Gzip,
|
|
|
+)
|
|
|
+_COMPRESSION_NAMES = {
|
|
|
+ None: 'Uncompressed',
|
|
|
+ grpc.Compression.NoCompression: 'NoCompression',
|
|
|
+ grpc.Compression.Deflate: 'DeflateCompression',
|
|
|
+ grpc.Compression.Gzip: 'GzipCompression',
|
|
|
+}
|
|
|
+
|
|
|
+_TEST_OPTIONS = {
|
|
|
+ 'client_streaming': (True, False),
|
|
|
+ 'server_streaming': (True, False),
|
|
|
+ 'channel_compression': _COMPRESSION_METHODS,
|
|
|
+ 'multicallable_compression': _COMPRESSION_METHODS,
|
|
|
+ 'server_compression': _COMPRESSION_METHODS,
|
|
|
+ 'server_call_compression': _COMPRESSION_METHODS,
|
|
|
+}
|
|
|
+
|
|
|
+
|
|
|
+def _make_handle_unary_unary(pre_response_callback):
|
|
|
+
|
|
|
+ def _handle_unary(request, servicer_context):
|
|
|
+ if pre_response_callback:
|
|
|
+ pre_response_callback(request, servicer_context)
|
|
|
+ return request
|
|
|
+
|
|
|
+ return _handle_unary
|
|
|
+
|
|
|
+
|
|
|
+def _make_handle_unary_stream(pre_response_callback):
|
|
|
+
|
|
|
+ def _handle_unary_stream(request, servicer_context):
|
|
|
+ if pre_response_callback:
|
|
|
+ pre_response_callback(request, servicer_context)
|
|
|
+ for _ in range(_STREAM_LENGTH):
|
|
|
+ yield request
|
|
|
+
|
|
|
+ return _handle_unary_stream
|
|
|
+
|
|
|
+
|
|
|
+def _make_handle_stream_unary(pre_response_callback):
|
|
|
+
|
|
|
+ def _handle_stream_unary(request_iterator, servicer_context):
|
|
|
+ if pre_response_callback:
|
|
|
+ pre_response_callback(request_iterator, servicer_context)
|
|
|
+ response = None
|
|
|
+ for request in request_iterator:
|
|
|
+ if not response:
|
|
|
+ response = request
|
|
|
+ return response
|
|
|
|
|
|
-def handle_unary(request, servicer_context):
|
|
|
- servicer_context.send_initial_metadata([('grpc-internal-encoding-request',
|
|
|
- 'gzip')])
|
|
|
- return request
|
|
|
+ return _handle_stream_unary
|
|
|
|
|
|
|
|
|
-def handle_stream(request_iterator, servicer_context):
|
|
|
- # TODO(issue:#6891) We should be able to remove this loop,
|
|
|
- # and replace with return; yield
|
|
|
- servicer_context.send_initial_metadata([('grpc-internal-encoding-request',
|
|
|
- 'gzip')])
|
|
|
- for request in request_iterator:
|
|
|
- yield request
|
|
|
+def _make_handle_stream_stream(pre_response_callback):
|
|
|
+
|
|
|
+ def _handle_stream(request_iterator, servicer_context):
|
|
|
+ # TODO(issue:#6891) We should be able to remove this loop,
|
|
|
+ # and replace with return; yield
|
|
|
+ for request in request_iterator:
|
|
|
+ if pre_response_callback:
|
|
|
+ pre_response_callback(request, servicer_context)
|
|
|
+ yield request
|
|
|
+
|
|
|
+ return _handle_stream
|
|
|
+
|
|
|
+
|
|
|
+def set_call_compression(compression_method, request_or_iterator,
|
|
|
+ servicer_context):
|
|
|
+ del request_or_iterator
|
|
|
+ servicer_context.set_compression(compression_method)
|
|
|
+
|
|
|
+
|
|
|
+def disable_next_compression(request, servicer_context):
|
|
|
+ del request
|
|
|
+ servicer_context.disable_next_message_compression()
|
|
|
+
|
|
|
+
|
|
|
+def disable_first_compression(request, servicer_context):
|
|
|
+ if int(request.decode('ascii')) == 0:
|
|
|
+ servicer_context.disable_next_message_compression()
|
|
|
|
|
|
|
|
|
class _MethodHandler(grpc.RpcMethodHandler):
|
|
|
|
|
|
- def __init__(self, request_streaming, response_streaming):
|
|
|
+ def __init__(self, request_streaming, response_streaming,
|
|
|
+ pre_response_callback):
|
|
|
self.request_streaming = request_streaming
|
|
|
self.response_streaming = response_streaming
|
|
|
self.request_deserializer = None
|
|
|
@@ -52,75 +145,239 @@ class _MethodHandler(grpc.RpcMethodHandler):
|
|
|
self.unary_stream = None
|
|
|
self.stream_unary = None
|
|
|
self.stream_stream = None
|
|
|
+
|
|
|
if self.request_streaming and self.response_streaming:
|
|
|
- self.stream_stream = handle_stream
|
|
|
+ self.stream_stream = _make_handle_stream_stream(
|
|
|
+ pre_response_callback)
|
|
|
elif not self.request_streaming and not self.response_streaming:
|
|
|
- self.unary_unary = handle_unary
|
|
|
+ self.unary_unary = _make_handle_unary_unary(pre_response_callback)
|
|
|
+ elif not self.request_streaming and self.response_streaming:
|
|
|
+ self.unary_stream = _make_handle_unary_stream(pre_response_callback)
|
|
|
+ else:
|
|
|
+ self.stream_unary = _make_handle_stream_unary(pre_response_callback)
|
|
|
|
|
|
|
|
|
class _GenericHandler(grpc.GenericRpcHandler):
|
|
|
|
|
|
+ def __init__(self, pre_response_callback):
|
|
|
+ self._pre_response_callback = pre_response_callback
|
|
|
+
|
|
|
def service(self, handler_call_details):
|
|
|
if handler_call_details.method == _UNARY_UNARY:
|
|
|
- return _MethodHandler(False, False)
|
|
|
+ return _MethodHandler(False, False, self._pre_response_callback)
|
|
|
+ elif handler_call_details.method == _UNARY_STREAM:
|
|
|
+ return _MethodHandler(False, True, self._pre_response_callback)
|
|
|
+ elif handler_call_details.method == _STREAM_UNARY:
|
|
|
+ return _MethodHandler(True, False, self._pre_response_callback)
|
|
|
elif handler_call_details.method == _STREAM_STREAM:
|
|
|
- return _MethodHandler(True, True)
|
|
|
+ return _MethodHandler(True, True, self._pre_response_callback)
|
|
|
else:
|
|
|
return None
|
|
|
|
|
|
|
|
|
+@contextlib.contextmanager
|
|
|
+def _instrumented_client_server_pair(channel_kwargs, server_kwargs,
|
|
|
+ server_handler):
|
|
|
+ server = grpc.server(futures.ThreadPoolExecutor(), **server_kwargs)
|
|
|
+ server.add_generic_rpc_handlers((server_handler,))
|
|
|
+ server_port = server.add_insecure_port('{}:0'.format(_HOST))
|
|
|
+ server.start()
|
|
|
+ with _tcp_proxy.TcpProxy(_HOST, _HOST, server_port) as proxy:
|
|
|
+ proxy_port = proxy.get_port()
|
|
|
+ with grpc.insecure_channel('{}:{}'.format(_HOST, proxy_port),
|
|
|
+ **channel_kwargs) as client_channel:
|
|
|
+ try:
|
|
|
+ yield client_channel, proxy, server
|
|
|
+ finally:
|
|
|
+ server.stop(None)
|
|
|
+
|
|
|
+
|
|
|
+def _get_byte_counts(channel_kwargs, multicallable_kwargs, client_function,
|
|
|
+ server_kwargs, server_handler, message):
|
|
|
+ with _instrumented_client_server_pair(channel_kwargs, server_kwargs,
|
|
|
+ server_handler) as pipeline:
|
|
|
+ client_channel, proxy, server = pipeline
|
|
|
+ client_function(client_channel, multicallable_kwargs, message)
|
|
|
+ return proxy.get_byte_count()
|
|
|
+
|
|
|
+
|
|
|
+def _get_compression_ratios(client_function, first_channel_kwargs,
|
|
|
+ first_multicallable_kwargs, first_server_kwargs,
|
|
|
+ first_server_handler, second_channel_kwargs,
|
|
|
+ second_multicallable_kwargs, second_server_kwargs,
|
|
|
+ second_server_handler, message):
|
|
|
+ try:
|
|
|
+ # This test requires the byte length of each connection to be deterministic. As
|
|
|
+ # it turns out, flow control puts bytes on the wire in a nondeterministic
|
|
|
+ # manner. We disable it here in order to measure compression ratios
|
|
|
+ # deterministically.
|
|
|
+ os.environ['GRPC_EXPERIMENTAL_DISABLE_FLOW_CONTROL'] = 'true'
|
|
|
+ first_bytes_sent, first_bytes_received = _get_byte_counts(
|
|
|
+ first_channel_kwargs, first_multicallable_kwargs, client_function,
|
|
|
+ first_server_kwargs, first_server_handler, message)
|
|
|
+ second_bytes_sent, second_bytes_received = _get_byte_counts(
|
|
|
+ second_channel_kwargs, second_multicallable_kwargs, client_function,
|
|
|
+ second_server_kwargs, second_server_handler, message)
|
|
|
+ return ((
|
|
|
+ second_bytes_sent - first_bytes_sent) / float(first_bytes_sent),
|
|
|
+ (second_bytes_received - first_bytes_received) /
|
|
|
+ float(first_bytes_received))
|
|
|
+ finally:
|
|
|
+ del os.environ['GRPC_EXPERIMENTAL_DISABLE_FLOW_CONTROL']
|
|
|
+
|
|
|
+
|
|
|
+def _unary_unary_client(channel, multicallable_kwargs, message):
|
|
|
+ multi_callable = channel.unary_unary(_UNARY_UNARY)
|
|
|
+ response = multi_callable(message, **multicallable_kwargs)
|
|
|
+ if response != message:
|
|
|
+ raise RuntimeError("Request '{}' != Response '{}'".format(
|
|
|
+ message, response))
|
|
|
+
|
|
|
+
|
|
|
+def _unary_stream_client(channel, multicallable_kwargs, message):
|
|
|
+ multi_callable = channel.unary_stream(_UNARY_STREAM)
|
|
|
+ response_iterator = multi_callable(message, **multicallable_kwargs)
|
|
|
+ for response in response_iterator:
|
|
|
+ if response != message:
|
|
|
+ raise RuntimeError("Request '{}' != Response '{}'".format(
|
|
|
+ message, response))
|
|
|
+
|
|
|
+
|
|
|
+def _stream_unary_client(channel, multicallable_kwargs, message):
|
|
|
+ multi_callable = channel.stream_unary(_STREAM_UNARY)
|
|
|
+ requests = (_REQUEST for _ in range(_STREAM_LENGTH))
|
|
|
+ response = multi_callable(requests, **multicallable_kwargs)
|
|
|
+ if response != message:
|
|
|
+ raise RuntimeError("Request '{}' != Response '{}'".format(
|
|
|
+ message, response))
|
|
|
+
|
|
|
+
|
|
|
+def _stream_stream_client(channel, multicallable_kwargs, message):
|
|
|
+ multi_callable = channel.stream_stream(_STREAM_STREAM)
|
|
|
+ request_prefix = str(0).encode('ascii') * 100
|
|
|
+ requests = (
|
|
|
+ request_prefix + str(i).encode('ascii') for i in range(_STREAM_LENGTH))
|
|
|
+ response_iterator = multi_callable(requests, **multicallable_kwargs)
|
|
|
+ for i, response in enumerate(response_iterator):
|
|
|
+ if int(response.decode('ascii')) != i:
|
|
|
+ raise RuntimeError("Request '{}' != Response '{}'".format(
|
|
|
+ i, response))
|
|
|
+
|
|
|
+
|
|
|
class CompressionTest(unittest.TestCase):
|
|
|
|
|
|
- def setUp(self):
|
|
|
- self._server = test_common.test_server()
|
|
|
- self._server.add_generic_rpc_handlers((_GenericHandler(),))
|
|
|
- self._port = self._server.add_insecure_port('[::]:0')
|
|
|
- self._server.start()
|
|
|
-
|
|
|
- def tearDown(self):
|
|
|
- self._server.stop(None)
|
|
|
-
|
|
|
- def testUnary(self):
|
|
|
- request = b'\x00' * 100
|
|
|
-
|
|
|
- # Client -> server compressed through default client channel compression
|
|
|
- # settings. Server -> client compressed via server-side metadata setting.
|
|
|
- # TODO(https://github.com/grpc/grpc/issues/4078): replace the "1" integer
|
|
|
- # literal with proper use of the public API.
|
|
|
- compressed_channel = grpc.insecure_channel(
|
|
|
- 'localhost:%d' % self._port,
|
|
|
- options=[('grpc.default_compression_algorithm', 1)])
|
|
|
- multi_callable = compressed_channel.unary_unary(_UNARY_UNARY)
|
|
|
- response = multi_callable(request)
|
|
|
- self.assertEqual(request, response)
|
|
|
-
|
|
|
- # Client -> server compressed through client metadata setting. Server ->
|
|
|
- # client compressed via server-side metadata setting.
|
|
|
- # TODO(https://github.com/grpc/grpc/issues/4078): replace the "0" integer
|
|
|
- # literal with proper use of the public API.
|
|
|
- uncompressed_channel = grpc.insecure_channel(
|
|
|
- 'localhost:%d' % self._port,
|
|
|
- options=[('grpc.default_compression_algorithm', 0)])
|
|
|
- multi_callable = compressed_channel.unary_unary(_UNARY_UNARY)
|
|
|
- response = multi_callable(
|
|
|
- request, metadata=[('grpc-internal-encoding-request', 'gzip')])
|
|
|
- self.assertEqual(request, response)
|
|
|
- compressed_channel.close()
|
|
|
-
|
|
|
- def testStreaming(self):
|
|
|
- request = b'\x00' * 100
|
|
|
-
|
|
|
- # TODO(https://github.com/grpc/grpc/issues/4078): replace the "1" integer
|
|
|
- # literal with proper use of the public API.
|
|
|
- compressed_channel = grpc.insecure_channel(
|
|
|
- 'localhost:%d' % self._port,
|
|
|
- options=[('grpc.default_compression_algorithm', 1)])
|
|
|
- multi_callable = compressed_channel.stream_stream(_STREAM_STREAM)
|
|
|
- call = multi_callable(iter([request] * test_constants.STREAM_LENGTH))
|
|
|
- for response in call:
|
|
|
- self.assertEqual(request, response)
|
|
|
- compressed_channel.close()
|
|
|
+ def assertCompressed(self, compression_ratio):
|
|
|
+ self.assertLess(
|
|
|
+ compression_ratio,
|
|
|
+ -1.0 * _COMPRESSION_RATIO_THRESHOLD,
|
|
|
+ msg='Actual compression ratio: {}'.format(compression_ratio))
|
|
|
+
|
|
|
+ def assertNotCompressed(self, compression_ratio):
|
|
|
+ self.assertGreaterEqual(
|
|
|
+ compression_ratio,
|
|
|
+ -1.0 * _COMPRESSION_RATIO_THRESHOLD,
|
|
|
+ msg='Actual compession ratio: {}'.format(compression_ratio))
|
|
|
+
|
|
|
+ def assertConfigurationCompressed(
|
|
|
+ self, client_streaming, server_streaming, channel_compression,
|
|
|
+ multicallable_compression, server_compression,
|
|
|
+ server_call_compression):
|
|
|
+ client_side_compressed = channel_compression or multicallable_compression
|
|
|
+ server_side_compressed = server_compression or server_call_compression
|
|
|
+ channel_kwargs = {
|
|
|
+ 'compression': channel_compression,
|
|
|
+ } if channel_compression else {}
|
|
|
+ multicallable_kwargs = {
|
|
|
+ 'compression': multicallable_compression,
|
|
|
+ } if multicallable_compression else {}
|
|
|
+
|
|
|
+ client_function = None
|
|
|
+ if not client_streaming and not server_streaming:
|
|
|
+ client_function = _unary_unary_client
|
|
|
+ elif not client_streaming and server_streaming:
|
|
|
+ client_function = _unary_stream_client
|
|
|
+ elif client_streaming and not server_streaming:
|
|
|
+ client_function = _stream_unary_client
|
|
|
+ else:
|
|
|
+ client_function = _stream_stream_client
|
|
|
+
|
|
|
+ server_kwargs = {
|
|
|
+ 'compression': server_compression,
|
|
|
+ } if server_compression else {}
|
|
|
+ server_handler = _GenericHandler(
|
|
|
+ functools.partial(set_call_compression, grpc.Compression.Gzip)
|
|
|
+ ) if server_call_compression else _GenericHandler(None)
|
|
|
+ sent_ratio, received_ratio = _get_compression_ratios(
|
|
|
+ client_function, {}, {}, {}, _GenericHandler(None), channel_kwargs,
|
|
|
+ multicallable_kwargs, server_kwargs, server_handler, _REQUEST)
|
|
|
+
|
|
|
+ if client_side_compressed:
|
|
|
+ self.assertCompressed(sent_ratio)
|
|
|
+ else:
|
|
|
+ self.assertNotCompressed(sent_ratio)
|
|
|
+
|
|
|
+ if server_side_compressed:
|
|
|
+ self.assertCompressed(received_ratio)
|
|
|
+ else:
|
|
|
+ self.assertNotCompressed(received_ratio)
|
|
|
+
|
|
|
+ def testDisableNextCompressionStreaming(self):
|
|
|
+ server_kwargs = {
|
|
|
+ 'compression': grpc.Compression.Deflate,
|
|
|
+ }
|
|
|
+ _, received_ratio = _get_compression_ratios(
|
|
|
+ _stream_stream_client, {}, {}, {}, _GenericHandler(None), {}, {},
|
|
|
+ server_kwargs, _GenericHandler(disable_next_compression), _REQUEST)
|
|
|
+ self.assertNotCompressed(received_ratio)
|
|
|
+
|
|
|
+ def testDisableNextCompressionStreamingResets(self):
|
|
|
+ server_kwargs = {
|
|
|
+ 'compression': grpc.Compression.Deflate,
|
|
|
+ }
|
|
|
+ _, received_ratio = _get_compression_ratios(
|
|
|
+ _stream_stream_client, {}, {}, {}, _GenericHandler(None), {}, {},
|
|
|
+ server_kwargs, _GenericHandler(disable_first_compression), _REQUEST)
|
|
|
+ self.assertCompressed(received_ratio)
|
|
|
+
|
|
|
+
|
|
|
+def _get_compression_str(name, value):
|
|
|
+ return '{}{}'.format(name, _COMPRESSION_NAMES[value])
|
|
|
+
|
|
|
+
|
|
|
+def _get_compression_test_name(client_streaming, server_streaming,
|
|
|
+ channel_compression, multicallable_compression,
|
|
|
+ server_compression, server_call_compression):
|
|
|
+ client_arity = 'Stream' if client_streaming else 'Unary'
|
|
|
+ server_arity = 'Stream' if server_streaming else 'Unary'
|
|
|
+ arity = '{}{}'.format(client_arity, server_arity)
|
|
|
+ channel_compression_str = _get_compression_str('Channel',
|
|
|
+ channel_compression)
|
|
|
+ multicallable_compression_str = _get_compression_str(
|
|
|
+ 'Multicallable', multicallable_compression)
|
|
|
+ server_compression_str = _get_compression_str('Server', server_compression)
|
|
|
+ server_call_compression_str = _get_compression_str('ServerCall',
|
|
|
+ server_call_compression)
|
|
|
+ return 'test{}{}{}{}{}'.format(
|
|
|
+ arity, channel_compression_str, multicallable_compression_str,
|
|
|
+ server_compression_str, server_call_compression_str)
|
|
|
+
|
|
|
+
|
|
|
+def _test_options():
|
|
|
+ for test_parameters in itertools.product(*_TEST_OPTIONS.values()):
|
|
|
+ yield dict(zip(_TEST_OPTIONS.keys(), test_parameters))
|
|
|
+
|
|
|
+
|
|
|
+for options in _test_options():
|
|
|
+
|
|
|
+ def test_compression(**kwargs):
|
|
|
+
|
|
|
+ def _test_compression(self):
|
|
|
+ self.assertConfigurationCompressed(**kwargs)
|
|
|
+
|
|
|
+ return _test_compression
|
|
|
|
|
|
+ setattr(CompressionTest, _get_compression_test_name(**options),
|
|
|
+ test_compression(**options))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
logging.basicConfig()
|
Richard Belleville