src/python/grpcio_tests/tests/stress/client.py

# Copyright 2016 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Entry point for running stress tests."""

import argparse
from concurrent import futures
import threading

import grpc
from six.moves import queue

from src.proto.grpc.testing import metrics_pb2_grpc
from src.proto.grpc.testing import test_pb2_grpc
from tests.interop import methods
from tests.interop import resources
from tests.qps import histogram
from tests.stress import metrics_server
from tests.stress import test_runner


def _args():
    parser = argparse.ArgumentParser(
        description='gRPC Python stress test client')
    parser.add_argument(
        '--server_addresses',
        help='comma separated list of hostname:port to run servers on',
        default='localhost:8080',
        type=str)
    parser.add_argument(
        '--test_cases',
        help='comma separated list of testcase:weighting of tests to run',
        default='large_unary:100',
        type=str)
    parser.add_argument('--test_duration_secs',
                        help='number of seconds to run the stress test',
                        default=-1,
                        type=int)
    parser.add_argument('--num_channels_per_server',
                        help='number of channels per server',
                        default=1,
                        type=int)
    parser.add_argument('--num_stubs_per_channel',
                        help='number of stubs to create per channel',
                        default=1,
                        type=int)
    parser.add_argument('--metrics_port',
                        help='the port to listen for metrics requests on',
                        default=8081,
                        type=int)
    parser.add_argument(
        '--use_test_ca',
        help='Whether to use our fake CA. Requires --use_tls=true',
        default=False,
        type=bool)
    parser.add_argument('--use_tls',
                        help='Whether to use TLS',
                        default=False,
                        type=bool)
    parser.add_argument('--server_host_override',
                        help='the server host to which to claim to connect',
                        type=str)
    return parser.parse_args()


def _test_case_from_arg(test_case_arg):
    for test_case in methods.TestCase:
        if test_case_arg == test_case.value:
            return test_case
    else:
        raise ValueError('No test case {}!'.format(test_case_arg))


def _parse_weighted_test_cases(test_case_args):
    weighted_test_cases = {}
    for test_case_arg in test_case_args.split(','):
        name, weight = test_case_arg.split(':', 1)
        test_case = _test_case_from_arg(name)
        weighted_test_cases[test_case] = int(weight)
    return weighted_test_cases


def _get_channel(target, args):
    if args.use_tls:
        if args.use_test_ca:
            root_certificates = resources.test_root_certificates()
        else:
            root_certificates = None  # will load default roots.
        channel_credentials = grpc.ssl_channel_credentials(
            root_certificates=root_certificates)
        options = ((
            'grpc.ssl_target_name_override',
            args.server_host_override,
        ),)
        channel = grpc.secure_channel(target,
                                      channel_credentials,
                                      options=options)
    else:
        channel = grpc.insecure_channel(target)

    # waits for the channel to be ready before we start sending messages
    grpc.channel_ready_future(channel).result()
    return channel


def run_test(args):
    test_cases = _parse_weighted_test_cases(args.test_cases)
    test_server_targets = args.server_addresses.split(',')
    # Propagate any client exceptions with a queue
    exception_queue = queue.Queue()
    stop_event = threading.Event()
    hist = histogram.Histogram(1, 1)
    runners = []

    server = grpc.server(futures.ThreadPoolExecutor(max_workers=25))
    metrics_pb2_grpc.add_MetricsServiceServicer_to_server(
        metrics_server.MetricsServer(hist), server)
    server.add_insecure_port('[::]:{}'.format(args.metrics_port))
    server.start()

    for test_server_target in test_server_targets:
        for _ in range(args.num_channels_per_server):
            channel = _get_channel(test_server_target, args)
            for _ in range(args.num_stubs_per_channel):
                stub = test_pb2_grpc.TestServiceStub(channel)
                runner = test_runner.TestRunner(stub, test_cases, hist,
                                                exception_queue, stop_event)
                runners.append(runner)

    for runner in runners:
        runner.start()
    try:
        timeout_secs = args.test_duration_secs
        if timeout_secs < 0:
            timeout_secs = None
        raise exception_queue.get(block=True, timeout=timeout_secs)
    except queue.Empty:
        # No exceptions thrown, success
        pass
    finally:
        stop_event.set()
        for runner in runners:
            runner.join()
        runner = None
        server.stop(None)


if __name__ == '__main__':
    run_test(_args())