1#!/usr/bin/env python3
2import argparse
3from concurrent.futures import ThreadPoolExecutor
4import http.cookiejar
5import importlib
6import json
7from pathlib import Path
8import ssl
9import subprocess
10import sys
11import time
12import urllib.parse
13import urllib.request
14import uuid
15
16
17def main():
18 parser = argparse.ArgumentParser()
19 parser.add_argument('--clients', type=int, default=24)
20 parser.add_argument('--requests', type=int, default=1200)
21 parser.add_argument('--idle-seconds', type=int, default=0)
22 parser.add_argument('--output', type=Path, required=True)
23 args = parser.parse_args()
24 assert 1 <= args.clients <= 64 and 12 <= args.requests <= 50000
25 assert 0 <= args.idle_seconds <= 600
26 for unit in ['studio-dashboard-unit-test', 'studio-host-unit-test']:
27 subprocess.run(['systemctl', 'is-active', '--quiet', unit], check=True)
28 repo = Path(__file__).resolve().parent.parent
29 sys.path.insert(0, str(repo / 'service/keycloak'))
30 from api import Keycloak, LoopbackHTTPS
31 Page = importlib.import_module('dashboard-shale-test').Page
32 secret = importlib.import_module('dashboard-run').secret
33 origin = 'https://globe.studio.test'
34 context = ssl.create_default_context(cafile='/var/lib/studio/ca-bundle.crt')
35
36 class TLS(urllib.request.HTTPSHandler):
37 def https_open(self, request):
38 target = urllib.parse.urlsplit(request.full_url)
39 assert target.scheme == 'https' and target.netloc in ('globe.studio.test', 'keycloak.studio.test')
40 return self.do_open(LoopbackHTTPS, request, context=context)
41
42 class Redirect(urllib.request.HTTPRedirectHandler):
43 def redirect_request(self, request, response, code, message, headers, target):
44 parts = urllib.parse.urlsplit(target)
45 assert parts.scheme == 'https' and parts.netloc in ('globe.studio.test', 'keycloak.studio.test')
46 return super().redirect_request(request, response, code, message, headers, target)
47
48 class NoRedirect(urllib.request.HTTPRedirectHandler):
49 def redirect_request(self, *args):
50 return None
51
52 keycloak = Keycloak('keycloak.studio.test', secret('get', 'keycloak', 'password'), attempts=1)
53 name = 'dashboard-load-' + uuid.uuid4().hex
54 password = uuid.uuid4().hex + 'A1!'
55 user = None
56 try:
57 role = keycloak.request('/admin/realms/master/roles/infra-admin')
58 user = keycloak.request('/admin/realms/master/users', 'POST', {
59 'username': name, 'enabled': True, 'requiredActions': [],
60 'credentials': [{'type': 'password', 'value': password, 'temporary': False}],
61 }, full=True)['id']
62 keycloak.request('/admin/realms/master/users/' + user + '/role-mappings/realm', 'POST', [role])
63 cookies = http.cookiejar.CookieJar()
64 login = urllib.request.build_opener(TLS(), Redirect(), urllib.request.HTTPCookieProcessor(cookies))
65 with login.open(origin + '/snow.oauth2/start?' + urllib.parse.urlencode({'rd': origin + '/'}), timeout=30) as response:
66 target = response.url
67 body = response.read(4 * 1024 * 1024 + 1).decode()
68 assert urllib.parse.urlsplit(target).netloc == 'keycloak.studio.test'
69 forms = [form for form in Page(body).forms if any(field.get('name') == 'password' for field in form['fields'])]
70 assert len(forms) == 1
71 action = urllib.parse.urljoin(target, forms[0]['action'])
72 action_url = urllib.parse.urlsplit(action)
73 assert action_url.scheme == 'https' and action_url.netloc == 'keycloak.studio.test'
74 fields = {field['name']: field.get('value', '') for field in forms[0]['fields'] if field.get('name')}
75 fields.update(username=name, password=password)
76 with login.open(urllib.request.Request(action, data=urllib.parse.urlencode(fields).encode(),
77 headers={'Content-Type': 'application/x-www-form-urlencoded', 'Origin': 'https://keycloak.studio.test'}), timeout=30) as response:
78 assert response.status == 200 and response.url == origin + '/'
79 response.read(4 * 1024 * 1024 + 1)
80 cookie_request = urllib.request.Request(origin + '/api/me')
81 cookies.add_cookie_header(cookie_request)
82 cookie = cookie_request.get_header('Cookie')
83 assert cookie
84 paths = ['/api/me', '/api/host', '/api/services', '/api/launcher', '/api/status', '/api/storage',
85 '/api/mcp', '/api/media/list', *['/api/metrics/' + metric + '?range=3600'
86 for metric in ['host.cpu', 'host.memory', 'service.cpu', 'service.memory']]]
87
88 def request(path):
89 start = time.monotonic()
90 client = urllib.request.build_opener(TLS(), NoRedirect())
91 with client.open(urllib.request.Request(origin + path, headers={'Cookie': cookie}), timeout=30) as response:
92 assert response.status == 200 and response.headers.get_content_type() == 'application/json'
93 body = response.read(8 * 1024 * 1024 + 1)
94 assert len(body) <= 8 * 1024 * 1024
95 value = json.loads(body)
96 if path == '/api/me':
97 assert value['name'] == name
98 if path == '/api/mcp':
99 assert {catalog['endpoint'] for catalog in value['catalogs']} == {origin + '/mcp/' + catalog
100 for catalog in ['observability', 'agents', 'shale']}
101 return path, (time.monotonic() - start) * 1000
102
103 for path in paths:
104 request(path)
105 start = time.monotonic()
106 with ThreadPoolExecutor(max_workers=args.clients) as clients:
107 results = list(clients.map(request, [paths[i % len(paths)] for i in range(args.requests)]))
108 report = {'public_https_and_real_sso': True, 'clients': args.clients, 'requests': len(results),
109 'seconds': round(time.monotonic() - start, 2), 'routes': {}}
110 for path in paths:
111 samples = sorted(ms for route, ms in results if route == path)
112 report['routes'][path] = {'requests': len(samples), 'p95_ms': round(samples[int(len(samples) * .95)], 2),
113 'max_ms': round(samples[-1], 2)}
114 print(json.dumps({'load_complete': report}), flush=True)
115 if args.idle_seconds:
116 units = ['studio-dashboard-unit-test.service', 'studio-host-unit-test.service']
117 before = {}
118 for unit in units:
119 stat = Path('/sys/fs/cgroup/system.slice') / unit / 'cpu.stat'
120 before[unit] = int(dict(line.split() for line in stat.read_text().splitlines())['usage_usec'])
121 start = time.monotonic()
122 time.sleep(args.idle_seconds)
123 elapsed = time.monotonic() - start
124 idle = {'seconds': round(elapsed, 2), 'percent_of_one_core': {}}
125 for unit in units:
126 stat = Path('/sys/fs/cgroup/system.slice') / unit / 'cpu.stat'
127 used = int(dict(line.split() for line in stat.read_text().splitlines())['usage_usec']) - before[unit]
128 idle['percent_of_one_core'][unit] = round(used / elapsed / 10000, 3)
129 request('/api/me')
130 report['idle'] = idle
131 finally:
132 if user:
133 keycloak = Keycloak('keycloak.studio.test', secret('get', 'keycloak', 'password'), attempts=1)
134 keycloak.request('/admin/realms/master/users/' + user, 'DELETE')
135 report['owned_sso_fixture_removed'] = True
136 args.output.write_text(json.dumps(report, indent=2) + '\n')
137 print(json.dumps(report), flush=True)
138
139
140if __name__ == '__main__':
141 main()