diff --git a/docs/testing-other-systems.rst b/docs/testing-other-systems.rst index 57a8a621ea..0f9dda3ba2 100644 --- a/docs/testing-other-systems.rst +++ b/docs/testing-other-systems.rst @@ -24,3 +24,24 @@ We can build a generic XML-RPC client, by wrapping :py:class:`xmlrpc.client.Serv .. literalinclude:: ../examples/custom_xmlrpc_client/xmlrpc_locustfile.py For more examples, see `locust-plugins `_ + +Example: writing a gRPC User/client +======================================= + +Similarly to the XML-RPC example, we can also load test a gRPC server. + +.. literalinclude:: ../examples/grpc/hello_server.py + +In this case, the gRPC stub methods can also be wrapped so that we can record the request stats. + +.. literalinclude:: ../examples/grpc/locustfile.py + +Note: In order to make the `grpcio` Python library gevent-compatible the following code needs to be executed before creating the gRPC channel. + +```python +import grpc.experimental.gevent as grpc_gevent +grpc_gevent.init_gevent() +``` + +Note: It is important to close the gRPC channel before stopping the User greenlet; otherwise Locust may not be able to stop executing. +This is due to an issue in `grpcio` (see `grpc#15880 `_). diff --git a/examples/grpc/hello.proto b/examples/grpc/hello.proto new file mode 100644 index 0000000000..ab0d21eacf --- /dev/null +++ b/examples/grpc/hello.proto @@ -0,0 +1,15 @@ +syntax = "proto3"; + +package locust.hello; + +service HelloService { + rpc SayHello (HelloRequest) returns (HelloResponse) {} +} + +message HelloRequest { + string name = 1; +} + +message HelloResponse { + string message = 1; +} diff --git a/examples/grpc/hello_pb2.py b/examples/grpc/hello_pb2.py new file mode 100644 index 0000000000..8bf74a3835 --- /dev/null +++ b/examples/grpc/hello_pb2.py @@ -0,0 +1,159 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: hello.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor.FileDescriptor( + name="hello.proto", + package="locust.hello", + syntax="proto3", + serialized_options=None, + create_key=_descriptor._internal_create_key, + serialized_pb=b'\n\x0bhello.proto\x12\x0clocust.hello"\x1c\n\x0cHelloRequest\x12\x0c\n\x04name\x18\x01 \x01(\t" \n\rHelloResponse\x12\x0f\n\x07message\x18\x01 \x01(\t2U\n\x0cHelloService\x12\x45\n\x08SayHello\x12\x1a.locust.hello.HelloRequest\x1a\x1b.locust.hello.HelloResponse"\x00\x62\x06proto3', +) + + +_HELLOREQUEST = _descriptor.Descriptor( + name="HelloRequest", + full_name="locust.hello.HelloRequest", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="name", + full_name="locust.hello.HelloRequest.name", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=29, + serialized_end=57, +) + + +_HELLORESPONSE = _descriptor.Descriptor( + name="HelloResponse", + full_name="locust.hello.HelloResponse", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="message", + full_name="locust.hello.HelloResponse.message", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax="proto3", + extension_ranges=[], + oneofs=[], + serialized_start=59, + serialized_end=91, +) + +DESCRIPTOR.message_types_by_name["HelloRequest"] = _HELLOREQUEST +DESCRIPTOR.message_types_by_name["HelloResponse"] = _HELLORESPONSE +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +HelloRequest = _reflection.GeneratedProtocolMessageType( + "HelloRequest", + (_message.Message,), + { + "DESCRIPTOR": _HELLOREQUEST, + "__module__": "hello_pb2" + # @@protoc_insertion_point(class_scope:locust.hello.HelloRequest) + }, +) +_sym_db.RegisterMessage(HelloRequest) + +HelloResponse = _reflection.GeneratedProtocolMessageType( + "HelloResponse", + (_message.Message,), + { + "DESCRIPTOR": _HELLORESPONSE, + "__module__": "hello_pb2" + # @@protoc_insertion_point(class_scope:locust.hello.HelloResponse) + }, +) +_sym_db.RegisterMessage(HelloResponse) + + +_HELLOSERVICE = _descriptor.ServiceDescriptor( + name="HelloService", + full_name="locust.hello.HelloService", + file=DESCRIPTOR, + index=0, + serialized_options=None, + create_key=_descriptor._internal_create_key, + serialized_start=93, + serialized_end=178, + methods=[ + _descriptor.MethodDescriptor( + name="SayHello", + full_name="locust.hello.HelloService.SayHello", + index=0, + containing_service=None, + input_type=_HELLOREQUEST, + output_type=_HELLORESPONSE, + serialized_options=None, + create_key=_descriptor._internal_create_key, + ), + ], +) +_sym_db.RegisterServiceDescriptor(_HELLOSERVICE) + +DESCRIPTOR.services_by_name["HelloService"] = _HELLOSERVICE + +# @@protoc_insertion_point(module_scope) diff --git a/examples/grpc/hello_pb2_grpc.py b/examples/grpc/hello_pb2_grpc.py new file mode 100644 index 0000000000..5f523bf0b0 --- /dev/null +++ b/examples/grpc/hello_pb2_grpc.py @@ -0,0 +1,77 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +import hello_pb2 as hello__pb2 + + +class HelloServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.SayHello = channel.unary_unary( + "/locust.hello.HelloService/SayHello", + request_serializer=hello__pb2.HelloRequest.SerializeToString, + response_deserializer=hello__pb2.HelloResponse.FromString, + ) + + +class HelloServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def SayHello(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + +def add_HelloServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + "SayHello": grpc.unary_unary_rpc_method_handler( + servicer.SayHello, + request_deserializer=hello__pb2.HelloRequest.FromString, + response_serializer=hello__pb2.HelloResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler("locust.hello.HelloService", rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + +# This class is part of an EXPERIMENTAL API. +class HelloService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def SayHello( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_unary( + request, + target, + "/locust.hello.HelloService/SayHello", + hello__pb2.HelloRequest.SerializeToString, + hello__pb2.HelloResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) diff --git a/examples/grpc/hello_server.py b/examples/grpc/hello_server.py new file mode 100644 index 0000000000..1307eaf651 --- /dev/null +++ b/examples/grpc/hello_server.py @@ -0,0 +1,24 @@ +import hello_pb2_grpc +import hello_pb2 +import grpc +from concurrent import futures +import logging +import time + +logger = logging.getLogger(__name__) + + +class HelloServiceServicer(hello_pb2_grpc.HelloServiceServicer): + def SayHello(self, request, context): + name = request.name + time.sleep(1) + return hello_pb2.HelloResponse(message=f"Hello from Locust, {name}!") + + +def start_server(): + server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + hello_pb2_grpc.add_HelloServiceServicer_to_server(HelloServiceServicer(), server) + server.add_insecure_port("localhost:50051") + server.start() + logger.info("gRPC server started") + server.wait_for_termination() diff --git a/examples/grpc/locustfile.py b/examples/grpc/locustfile.py new file mode 100644 index 0000000000..2a0dd83833 --- /dev/null +++ b/examples/grpc/locustfile.py @@ -0,0 +1,82 @@ +import grpc +import hello_pb2_grpc +import hello_pb2 +from locust import events, User, task +from locust.exception import LocustError +from locust.user.task import LOCUST_STATE_STOPPING +from hello_server import start_server +import gevent +import time + +# patch grpc so that it uses gevent instead of asyncio +import grpc.experimental.gevent as grpc_gevent + +grpc_gevent.init_gevent() + + +@events.init.add_listener +def run_grpc_server(environment, **_kwargs): + gevent.spawn(start_server) + + +class GrpcClient: + def __init__(self, stub): + self._stub_class = stub.__class__ + self._stub = stub + + def __getattr__(self, name): + func = self._stub_class.__getattribute__(self._stub, name) + + def wrapper(*args, **kwargs): + start_time = time.monotonic() + request_meta = { + "request_type": "grpc", + "name": name, + "response_length": 0, + "exception": None, + "context": None, + "response": None, + } + try: + request_meta["response"] = func(*args, **kwargs) + request_meta["response_length"] = len(request_meta["response"].message) + except grpc.RpcError as e: + request_meta["exception"] = e + request_meta["response_time"] = (time.monotonic() - start_time) * 1000 + events.request.fire(**request_meta) + return request_meta["response"] + + return wrapper + + +class GrpcUser(User): + abstract = True + + stub_class = None + + def __init__(self, environment): + super().__init__(environment) + for attr_value, attr_name in ((self.host, "host"), (self.stub_class, "stub_class")): + if attr_value is None: + raise LocustError(f"You must specify the {attr_name}.") + self._channel = grpc.insecure_channel(self.host) + self._channel_closed = False + stub = self.stub_class(self._channel) + self.client = GrpcClient(stub) + + def stop(self, force=False): + self._channel_closed = True + time.sleep(1) + self._channel.close() + super().stop(force=True) + + +class HelloGrpcUser(GrpcUser): + host = "localhost:50051" + stub_class = hello_pb2_grpc.HelloServiceStub + + @task + def sayHello(self): + if not self._channel_closed: + self.client.SayHello(hello_pb2.HelloRequest(name="Test")) + time.sleep(1) diff --git a/locust/runners.py b/locust/runners.py index f3f4684c37..0c38fc3825 100644 --- a/locust/runners.py +++ b/locust/runners.py @@ -250,10 +250,10 @@ def stop_users(self, user_count, stop_rate=None): # User called runner.quit(), so dont block waiting for killing to finish" user_to_stop._group.killone(user_to_stop._greenlet, block=False) elif self.environment.stop_timeout: - async_calls_to_stop.add(gevent.spawn_later(0, User.stop, user_to_stop, force=False)) + async_calls_to_stop.add(gevent.spawn_later(0, user_to_stop.stop, force=False)) stop_group.add(user_to_stop._greenlet) else: - async_calls_to_stop.add(gevent.spawn_later(0, User.stop, user_to_stop, force=True)) + async_calls_to_stop.add(gevent.spawn_later(0, user_to_stop.stop, force=True)) if to_stop: gevent.sleep(sleep_time) else: