1#!/usr/bin/env python3
2# Copyright 2021 The gRPC Authors
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16# This script generates a load test configuration template from a collection of
17# load test configurations.
18#
19# Configuration templates contain client and server configurations for multiple
20# languages, and may contain template substitution keys. These templates are
21# used to generate load test configurations by selecting clients and servers for
22# the required languages. The source files for template generation may be load
23# test configurations or load test configuration templates. Load test
24# configuration generation is performed by loadtest_config.py. See documentation
25# below:
26# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md
27
28import argparse
29import os
30import sys
31from typing import Any, Dict, Iterable, List, Mapping, Type
32
33import yaml
34
35sys.path.append(os.path.dirname(os.path.abspath(__file__)))
36import loadtest_config
37
38TEMPLATE_FILE_HEADER_COMMENT = """
39# Template generated from load test configurations by loadtest_template.py.
40#
41# Configuration templates contain client and server configurations for multiple
42# languages, and may contain template substitution keys. These templates are
43# used to generate load test configurations by selecting clients and servers for
44# the required languages. The source files for template generation may be load
45# test configurations or load test configuration templates. Load test
46# configuration generation is performed by loadtest_config.py. See documentation
47# below:
48# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md
49"""
50
51
52def insert_worker(worker: Dict[str, Any], workers: List[Dict[str,
53                                                             Any]]) -> None:
54    """Inserts client or server into a list, without inserting duplicates."""
55
56    def dump(w):
57        return yaml.dump(w, Dumper=yaml.SafeDumper, default_flow_style=False)
58
59    worker_str = dump(worker)
60    if any((worker_str == dump(w) for w in workers)):
61        return
62    workers.append(worker)
63
64
65def uniquify_workers(workermap: Dict[str, List[Dict[str, Any]]]) -> None:
66    """Name workers if there is more than one for the same map key."""
67    for workers in list(workermap.values()):
68        if len(workers) <= 1:
69            continue
70        for i, worker in enumerate(workers):
71            worker['name'] = str(i)
72
73
74def loadtest_template(
75        input_file_names: Iterable[str],
76        metadata: Mapping[str, Any],
77        inject_client_pool: bool,
78        inject_driver_image: bool,
79        inject_driver_pool: bool,
80        inject_server_pool: bool,
81        inject_big_query_table: bool,
82        inject_timeout_seconds: bool,
83        inject_ttl_seconds: bool) -> Dict[str, Any]:  # yapf: disable
84    """Generates the load test template."""
85    spec = dict()  # type: Dict[str, Any]
86    clientmap = dict()  # Dict[str, List[Dict[str, Any]]]
87    servermap = dict()  # Dict[Str, List[Dict[str, Any]]]
88    template = {
89        'apiVersion': 'e2etest.grpc.io/v1',
90        'kind': 'LoadTest',
91        'metadata': metadata,
92    }
93    for input_file_name in input_file_names:
94        with open(input_file_name) as f:
95            input_config = yaml.safe_load(f.read())
96
97            if input_config.get('apiVersion') != template['apiVersion']:
98                raise ValueError('Unexpected api version in file {}: {}'.format(
99                    input_file_name, input_config.get('apiVersion')))
100            if input_config.get('kind') != template['kind']:
101                raise ValueError('Unexpected kind in file {}: {}'.format(
102                    input_file_name, input_config.get('kind')))
103
104            for client in input_config['spec']['clients']:
105                del client['name']
106                if inject_client_pool:
107                    client['pool'] = '${client_pool}'
108                if client['language'] not in clientmap:
109                    clientmap[client['language']] = []
110                insert_worker(client, clientmap[client['language']])
111
112            for server in input_config['spec']['servers']:
113                del server['name']
114                if inject_server_pool:
115                    server['pool'] = '${server_pool}'
116                if server['language'] not in servermap:
117                    servermap[server['language']] = []
118                insert_worker(server, servermap[server['language']])
119
120            input_spec = input_config['spec']
121            del input_spec['clients']
122            del input_spec['servers']
123            del input_spec['scenariosJSON']
124            spec.update(input_config['spec'])
125
126    uniquify_workers(clientmap)
127    uniquify_workers(servermap)
128
129    spec.update({
130        'clients':
131            sum((clientmap[language] for language in sorted(clientmap)),
132                start=[]),
133        'servers':
134            sum((servermap[language] for language in sorted(servermap)),
135                start=[]),
136    })
137
138    if 'driver' not in spec:
139        spec['driver'] = {'language': 'cxx'}
140
141    driver = spec['driver']
142    if 'name' in driver:
143        del driver['name']
144    if inject_driver_image:
145        if 'run' not in driver:
146            driver['run'] = [{'name': 'main'}]
147        driver['run'][0]['image'] = '${driver_image}'
148    if inject_driver_pool:
149        driver['pool'] = '${driver_pool}'
150
151    if 'run' not in driver:
152        if inject_driver_pool:
153            raise ValueError('Cannot inject driver.pool: missing driver.run.')
154        del spec['driver']
155
156    if inject_big_query_table:
157        if 'results' not in spec:
158            spec['results'] = dict()
159        spec['results']['bigQueryTable'] = '${big_query_table}'
160    if inject_timeout_seconds:
161        spec['timeoutSeconds'] = '${timeout_seconds}'
162    if inject_ttl_seconds:
163        spec['ttlSeconds'] = '${ttl_seconds}'
164
165    template['spec'] = spec
166
167    return template
168
169
170def template_dumper(header_comment: str) -> Type[yaml.SafeDumper]:
171    """Returns a custom dumper to dump templates in the expected format."""
172
173    class TemplateDumper(yaml.SafeDumper):
174
175        def expect_stream_start(self):
176            super().expect_stream_start()
177            if isinstance(self.event, yaml.StreamStartEvent):
178                self.write_indent()
179                self.write_indicator(header_comment, need_whitespace=False)
180
181    def str_presenter(dumper, data):
182        if '\n' in data:
183            return dumper.represent_scalar('tag:yaml.org,2002:str',
184                                           data,
185                                           style='|')
186        return dumper.represent_scalar('tag:yaml.org,2002:str', data)
187
188    TemplateDumper.add_representer(str, str_presenter)
189
190    return TemplateDumper
191
192
193def main() -> None:
194    argp = argparse.ArgumentParser(
195        description='Creates a load test config generator template.',
196        fromfile_prefix_chars='@')
197    argp.add_argument('-i',
198                      '--inputs',
199                      action='extend',
200                      nargs='+',
201                      type=str,
202                      help='Input files.')
203    argp.add_argument('-o',
204                      '--output',
205                      type=str,
206                      help='Output file. Outputs to stdout if not set.')
207    argp.add_argument(
208        '--inject_client_pool',
209        action='store_true',
210        help='Set spec.client(s).pool values to \'${client_pool}\'.')
211    argp.add_argument(
212        '--inject_driver_image',
213        action='store_true',
214        help='Set spec.driver(s).image values to \'${driver_image}\'.')
215    argp.add_argument(
216        '--inject_driver_pool',
217        action='store_true',
218        help='Set spec.driver(s).pool values to \'${driver_pool}\'.')
219    argp.add_argument(
220        '--inject_server_pool',
221        action='store_true',
222        help='Set spec.server(s).pool values to \'${server_pool}\'.')
223    argp.add_argument(
224        '--inject_big_query_table',
225        action='store_true',
226        help='Set spec.results.bigQueryTable to \'${big_query_table}\'.')
227    argp.add_argument('--inject_timeout_seconds',
228                      action='store_true',
229                      help='Set spec.timeoutSeconds to \'${timeout_seconds}\'.')
230    argp.add_argument('--inject_ttl_seconds',
231                      action='store_true',
232                      help='Set timeout ')
233    argp.add_argument('-n',
234                      '--name',
235                      default='',
236                      type=str,
237                      help='metadata.name.')
238    argp.add_argument('-a',
239                      '--annotation',
240                      action='append',
241                      type=str,
242                      help='metadata.annotation(s), in the form key=value.',
243                      dest='annotations')
244    args = argp.parse_args()
245
246    annotations = loadtest_config.parse_key_value_args(args.annotations)
247
248    metadata = {'name': args.name}
249    if annotations:
250        metadata['annotations'] = annotations
251
252    template = loadtest_template(
253        input_file_names=args.inputs,
254        metadata=metadata,
255        inject_client_pool=args.inject_client_pool,
256        inject_driver_image=args.inject_driver_image,
257        inject_driver_pool=args.inject_driver_pool,
258        inject_server_pool=args.inject_server_pool,
259        inject_big_query_table=args.inject_big_query_table,
260        inject_timeout_seconds=args.inject_timeout_seconds,
261        inject_ttl_seconds=args.inject_ttl_seconds)
262
263    with open(args.output, 'w') if args.output else sys.stdout as f:
264        yaml.dump(template,
265                  stream=f,
266                  Dumper=template_dumper(TEMPLATE_FILE_HEADER_COMMENT.strip()),
267                  default_flow_style=False)
268
269
270if __name__ == '__main__':
271    main()
272