Locust

.github/workflows/ci.yml

@@ -14,6 +14,10 @@ jobs:
       - name: Check out the repo
         uses: actions/checkout@v3
 
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.11'
+
       - name: Build docker containers
         run: docker compose build
 
@@ -29,6 +33,30 @@ jobs:
       - name: Test client runs without errors
         run: docker compose run --rm client
 
+      - name: Run load test
+        run: |
+          python --version
+          pip install -r load-test/requirements.txt
+          python -m grpc_tools.protoc --proto_path=datastore/protobuf datastore.proto --python_out=load-test --grpc_python_out=load-test
+          cd load-test
+          locust --headless -u 5 -r 1 --run-time 60 --only-summary --csv store
+
+      - name: Archive load test artifacts
+        uses: actions/upload-artifact@v3
+        with:
+          name: performance
+          path: load-test/store_*.csv
+
+      - name: Print results
+        run: |
+          pip install csvkit
+          echo "## Stats" >> $GITHUB_STEP_SUMMARY
+          csvlook load-test/store_stats.csv >> $GITHUB_STEP_SUMMARY
+          echo "## Stats history" >> $GITHUB_STEP_SUMMARY
+          csvlook load-test/store_stats_history.csv >> $GITHUB_STEP_SUMMARY
+          echo "## Failures" >> $GITHUB_STEP_SUMMARY
+          csvlook load-test/store_failures.csv >> $GITHUB_STEP_SUMMARY
+
       - name: Cleanup
         if: always()
         run: docker compose down --volumes

load-test/grpc_user.py

@@ -0,0 +1,62 @@
+import time
+from typing import Any, Callable
+import grpc
+import grpc.experimental.gevent as grpc_gevent
+from grpc_interceptor import ClientInterceptor
+from locust import User
+from locust.exception import LocustError
+
+# patch grpc so that it uses gevent instead of asyncio
+grpc_gevent.init_gevent()
+
+
+class LocustInterceptor(ClientInterceptor):
+    def __init__(self, environment, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.env = environment
+
+    def intercept(
+        self,
+        method: Callable,
+        request_or_iterator: Any,
+        call_details: grpc.ClientCallDetails,
+    ):
+        response = None
+        exception = None
+        start_perf_counter = time.perf_counter()
+        response_length = 0
+        try:
+            response = method(request_or_iterator, call_details)
+            response_length = response.result().ByteSize()
+        except grpc.RpcError as e:
+            exception = e
+
+        self.env.events.request.fire(
+            request_type="grpc",
+            name=call_details.method,
+            response_time=(time.perf_counter() - start_perf_counter) * 1000,
+            response_length=response_length,
+            response=response,
+            context=None,
+            exception=exception,
+        )
+        return response
+
+
+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)
+        interceptor = LocustInterceptor(environment=environment)
+        self._channel = grpc.intercept_channel(self._channel, interceptor)
+
+        self.stub = self.stub_class(self._channel)
+

load-test/locustfile.py

@@ -0,0 +1,43 @@
+# Use the following command to generate the python protobuf stuff in the correct place (from the root of the repository)
+# python -m grpc_tools.protoc --proto_path=datastore/protobuf datastore.proto --python_out=load-test --grpc_python_out=load-test
+
+import random
+from datetime import datetime
+
+import grpc_user
+import datastore_pb2 as dstore
+import datastore_pb2_grpc as dstore_grpc
+from locust import task
+
+from google.protobuf.timestamp_pb2 import Timestamp
+
+
+class StoreGrpcUser(grpc_user.GrpcUser):
+    host = "localhost:50050"
+    stub_class = dstore_grpc.DatastoreStub
+
+    @task
+    def find_debilt_humidity(self):
+        ts_request = dstore.FindTSRequest(
+            station_ids=["06260"],
+            param_ids=["rh"]
+        )
+        ts_response = self.stub.FindTimeSeries(ts_request)
+        assert len(ts_response.tseries) == 1
+
+    @task
+    def get_data_random_timeserie(self):
+        ts_id = random.randint(1, 55*44)
+
+        from_time = Timestamp()
+        from_time.FromDatetime(datetime(2022, 12, 31))
+        to_time = Timestamp()
+        to_time.FromDatetime(datetime(2023, 11, 1))
+        request = dstore.GetObsRequest(
+            tsids=[ts_id],
+            fromtime=from_time,
+            totime=to_time,
+        )
+        response = self.stub.GetObservations(request)
+        assert len(response.tsobs[0].obs) == 144
+

load-test/requirements.in

@@ -0,0 +1,8 @@
+# Generate requirements.txt using:
+# pip-compile --upgrade --no-emit-index-url
+# Install using:
+# pip-sync
+
+grpcio-tools~=1.56
+grpc-interceptor~=0.15.3
+locust~=2.16

load-test/requirements.txt

@@ -0,0 +1,86 @@
+#
+# This file is autogenerated by pip-compile with Python 3.11
+# by the following command:
+#
+#    pip-compile --no-emit-index-url
+#
+blinker==1.6.2
+    # via flask
+brotli==1.1.0
+    # via geventhttpclient
+certifi==2023.7.22
+    # via
+    #   geventhttpclient
+    #   requests
+charset-normalizer==3.2.0
+    # via requests
+click==8.1.7
+    # via flask
+configargparse==1.7
+    # via locust
+flask==2.3.3
+    # via
+    #   flask-basicauth
+    #   flask-cors
+    #   locust
+flask-basicauth==0.2.0
+    # via locust
+flask-cors==4.0.0
+    # via locust
+gevent==23.9.0.post1
+    # via
+    #   geventhttpclient
+    #   locust
+geventhttpclient==2.0.10
+    # via locust
+greenlet==2.0.2
+    # via gevent
+grpc-interceptor==0.15.3
+    # via -r requirements.in
+grpcio==1.58.0
+    # via
+    #   grpc-interceptor
+    #   grpcio-tools
+grpcio-tools==1.58.0
+    # via -r requirements.in
+idna==3.4
+    # via requests
+itsdangerous==2.1.2
+    # via flask
+jinja2==3.1.2
+    # via flask
+locust==2.16.1
+    # via -r requirements.in
+markupsafe==2.1.3
+    # via
+    #   jinja2
+    #   werkzeug
+msgpack==1.0.5
+    # via locust
+protobuf==4.24.3
+    # via grpcio-tools
+psutil==5.9.5
+    # via locust
+pyzmq==25.1.1
+    # via locust
+requests==2.31.0
+    # via locust
+roundrobin==0.0.4
+    # via locust
+six==1.16.0
+    # via geventhttpclient
+typing-extensions==4.7.1
+    # via locust
+urllib3==2.0.4
+    # via requests
+werkzeug==2.3.7
+    # via
+    #   flask
+    #   locust
+zope-event==5.0
+    # via gevent
+zope-interface==6.0
+    # via gevent
+
+# The following packages are considered to be unsafe in a requirements file:
+# setuptools