[WIP] Ref #697 - Add profiling for whois and HTTP

docs/admins/configuration.rst

@@ -41,6 +41,7 @@ This sample shows most configuration options
         piddir: /var/run/
         user: irrd
         group: irrd
+        profiling_available: true
 
         access_lists:
             http_database_status:
@@ -220,6 +221,12 @@ General settings
   need for IRRd to bind to port 80 or 443.
   |br| **Default**: not defined, IRRd does not drop privileges.
   |br| **Change takes effect**: after full IRRd restart.
+* ``profiling_available``: whether to allow profiling output with the
+  ``!fprofile`` whois command or adding the ``profile=1`` GET parameter to
+  HTTP requests. Profiling is only available for CPython, so this setting
+  has no effect for PyPy deployments.
+  |br| **Default**: false.
+  |br| **Change takes effect**: after SIGHUP, on the next request/connection.
 
 
 Servers

irrd/conf/known_keys.py

@@ -6,6 +6,7 @@
 # Note that sources are checked separately,
 # and 'access_lists' is always permitted
 KNOWN_CONFIG_KEYS = DottedDict({
+    'profiling_available': {},
     'database_url': {},
     'database_readonly': {},
     'redis_url': {},

irrd/server/http/app.py

@@ -1,18 +1,21 @@
 import logging
 import os
+import platform
 import signal
 
 from ariadne.asgi import GraphQL
 from ariadne.asgi.handlers import GraphQLHTTPHandler
 from setproctitle import setproctitle
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import HTMLResponse
 from starlette.routing import Mount, Route, WebSocketRoute
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 # Relative imports are not allowed in this file
 from irrd import ENV_MAIN_PROCESS_PID
-from irrd.conf import config_init
+from irrd.conf import config_init, get_setting
 from irrd.server.graphql import ENV_UVICORN_WORKER_CONFIG_PATH
 from irrd.server.graphql.extensions import error_formatter, QueryMetadataExtension
 from irrd.server.graphql.schema_builder import build_executable_schema
@@ -28,6 +31,10 @@
 from irrd.utils.process_support import memory_trim
 
 
+if platform.python_implementation() == "CPython":
+    from pyinstrument import Profiler  # pragma: no cover
+
+
 logger = logging.getLogger(__name__)
 
 """
@@ -101,10 +108,23 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         memory_trim()
 
 
+class ProfileMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        profiling = platform.python_implementation() and get_setting('profiling_available') and request.query_params.get("profile", False)
+        if profiling:
+            profiler = Profiler(async_mode="enabled")
+            profiler.start()
+            await call_next(request)
+            profiler.stop()
+            return HTMLResponse(profiler.output_html())
+        else:
+            return await call_next(request)
+
+
 app = Starlette(
     debug=False,
     routes=routes,
     on_startup=[startup],
     on_shutdown=[shutdown],
-    middleware=[Middleware(MemoryTrimMiddleware)],
+    middleware=[Middleware(MemoryTrimMiddleware), Middleware(ProfileMiddleware)],
 )

irrd/server/whois/query_parser.py

@@ -1,4 +1,5 @@
 import logging
+import platform
 import re
 from typing import Optional
 
@@ -34,6 +35,7 @@ class WhoisQueryParser:
     so a single instance of this object should be created per session, with
     handle_query() being called for each individual query.
     """
+    profiling_enabled = False
 
     def __init__(self, client_ip: str, client_str: str, preloader: Preloader,
                  database_handler: DatabaseHandler) -> None:
@@ -123,6 +125,13 @@ def handle_irrd_command(self, full_command: str) -> WhoisQueryResponse:
             self.query_resolver.disable_out_of_scope_filter()
             result = 'Filtering out out-of-scope objects is disabled for !r and RIPE style ' \
                      'queries for the rest of this connection.'
+        elif all([
+            full_command.upper() == 'FPROFILE',
+            get_setting('profiling_available'),
+            platform.python_implementation() == "CPython"
+        ]):
+            self.profiling_enabled = True
+            result = 'Profiling output is enabled. Performance may be affected. '
         elif command == 'v':
             result = self.handle_irrd_version()
         elif command == 't':

irrd/server/whois/server.py

@@ -1,6 +1,7 @@
 import logging
 import multiprocessing as mp
 import os
+import platform
 import signal
 import socket
 import socketserver
@@ -19,6 +20,12 @@
 from irrd.storage.preload import Preloader
 from irrd.utils.process_support import memory_trim
 
+if platform.python_implementation() == "CPython":
+    from pyinstrument import Profiler  # pragma: no cover
+
+
+WHOIS_PROFILER_INTERVAL = 0.0001
+
 logger = logging.getLogger(__name__)
 mp.allow_connection_pickling()
 
@@ -213,10 +220,22 @@ def handle_query(self, query: str) -> bool:
             logger.debug(f'{self.client_str}: closed connection per request')
             return False
 
-        response = self.query_parser.handle_query(query)
-        response_bytes = response.generate_response().encode('utf-8')
+        profiler = None
+        if self.query_parser.profiling_enabled:   # pragma: no cover (CPython-only)
+            profiler = Profiler(interval=WHOIS_PROFILER_INTERVAL, async_mode="enabled")
+            profiler.start()
+            response = self.query_parser.handle_query(query)
+            response_bytes = response.generate_response().encode('utf-8')
+            profiler.stop()
+        else:
+            response = self.query_parser.handle_query(query)
+            response_bytes = response.generate_response().encode('utf-8')
+
         try:
             self.wfile.write(response_bytes)
+            if profiler:  # pragma: no cover (CPython-only)
+                profiler_output = profiler.output_text(show_all=True, unicode=True, color=True)
+                self.wfile.write(profiler_output.encode('utf-8'))
         except OSError:
             return False
 

irrd/server/whois/tests/test_query_parser.py

@@ -1,3 +1,4 @@
+import platform
 import uuid
 from unittest.mock import Mock
 
@@ -88,7 +89,7 @@
 
 
 @pytest.fixture()
-def prepare_parser(monkeypatch, config_override):
+def prepare_parser(monkeypatch):
     mock_query_resolver = Mock(spec=QueryResolver)
     mock_query_resolver.rpki_aware = False
     monkeypatch.setattr('irrd.server.whois.query_parser.QueryResolver',
@@ -828,6 +829,21 @@ def test_disable_filters(self, prepare_parser):
         assert response.result.startswith('Filtering out out-of-scope')
         mock_query_resolver.disable_out_of_scope_filter.assert_called_once_with()
 
+    @pytest.mark.skipif(platform.python_implementation() != "CPython", reason="requires CPython")
+    def test_profile(self, prepare_parser, config_override):
+        mock_query_resolver, mock_dh, parser = prepare_parser
+
+        assert not parser.profiling_enabled
+        parser.handle_query('!profile')
+        assert not parser.profiling_enabled
+
+        config_override({'profiling_available': True})
+        response = parser.handle_query('!fprofile')
+        assert parser.profiling_enabled
+        assert response.response_type == WhoisQueryResponseType.SUCCESS
+        assert response.mode == WhoisQueryResponseMode.IRRD
+        assert response.result.startswith('Profiling')
+
     def test_exception_handling(self, prepare_parser, caplog):
         mock_query_resolver, mock_dh, parser = prepare_parser
         mock_query_resolver.members_for_set = Mock(side_effect=Exception('test-error'))

irrd/server/whois/tests/test_server.py

@@ -1,3 +1,4 @@
+import platform
 import socket
 import time
 from io import BytesIO
@@ -66,6 +67,21 @@ def test_whois_request_worker_no_access_list(self, create_worker):
         assert request.close_called
         assert request.timeout_set == 5
 
+    @pytest.mark.skipif(platform.python_implementation() != "CPython", reason="requires CPython")
+    def test_whois_request_profile(self, create_worker, config_override):
+        config_override({'profiling_available': True})
+        worker, request = create_worker
+        # Empty query in first line should be ignored.
+        request.rfile.write(b'!!\n!fprofile\n!v\r\n')
+        request.rfile.seek(0)
+        worker.run(keep_running=False)
+
+        assert worker.client_str == '192.0.2.1:99999'
+        request.wfile.seek(0)
+        response = request.wfile.read()
+        assert b'IRRd -- version' in response
+        assert b'Recorded' in response
+
     def test_whois_request_worker_exception(self, create_worker, monkeypatch, caplog):
         monkeypatch.setattr('irrd.server.whois.server.WhoisQueryParser',
                             Mock(side_effect=OSError('expected')))

requirements.txt

@@ -26,6 +26,7 @@ psutil==5.9.4          # Process management
 asgiref==3.6.0         # ASGI utilities
 pydantic==1.10.2        # Input validation
 typing-extensions==4.4.0
+pyinstrument==4.4.0; platform_python_implementation == "CPython"    # Profiling requests
 
 # Database connections and management
 psycopg2-binary==2.9.5; platform_python_implementation == "CPython"