1#!/usr/bin/env python3
2#
3# Copyright 2022 gRPC authors.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import argparse
18import csv
19import glob
20import math
21import multiprocessing
22import os
23import pathlib
24import re
25import shutil
26import subprocess
27import sys
28
29sys.path.append(
30    os.path.join(os.path.dirname(sys.argv[0]), '..', '..', 'run_tests',
31                 'python_utils'))
32import check_on_pr
33
34argp = argparse.ArgumentParser(description='Perform diff on memory benchmarks')
35
36argp.add_argument('-d',
37                  '--diff_base',
38                  type=str,
39                  help='Commit or branch to compare the current one to')
40
41argp.add_argument('-j', '--jobs', type=int, default=multiprocessing.cpu_count())
42
43args = argp.parse_args()
44
45_INTERESTING = {
46    'call/client':
47        (rb'client call memory usage: ([0-9\.]+) bytes per call', float),
48    'call/server':
49        (rb'server call memory usage: ([0-9\.]+) bytes per call', float),
50    'channel/client':
51        (rb'client channel memory usage: ([0-9\.]+) bytes per channel', float),
52    'channel/server':
53        (rb'server channel memory usage: ([0-9\.]+) bytes per channel', float),
54}
55
56_SCENARIOS = {
57    'default': [],
58    'minstack': ['--scenario_config=minstack'],
59}
60
61_BENCHMARKS = {
62    'call': ['--benchmark_names=call', '--size=50000'],
63    'channel': ['--benchmark_names=channel', '--size=10000'],
64}
65
66
67def _run():
68    """Build with Bazel, then run, and extract interesting lines from the output."""
69    subprocess.check_call([
70        'tools/bazel', 'build', '-c', 'opt',
71        'test/core/memory_usage/memory_usage_test'
72    ])
73    ret = {}
74    for name, benchmark_args in _BENCHMARKS.items():
75        for scenario, extra_args in _SCENARIOS.items():
76            #TODO(chenancy) Remove when minstack is implemented for channel
77            if name == 'channel' and scenario == 'minstack':
78                continue
79            try:
80                output = subprocess.check_output([
81                    'bazel-bin/test/core/memory_usage/memory_usage_test',
82                ] + benchmark_args + extra_args)
83            except subprocess.CalledProcessError as e:
84                print('Error running benchmark:', e)
85                continue
86            for line in output.splitlines():
87                for key, (pattern, conversion) in _INTERESTING.items():
88                    m = re.match(pattern, line)
89                    if m:
90                        ret[scenario + ': ' + key] = conversion(m.group(1))
91    return ret
92
93
94cur = _run()
95old = None
96
97if args.diff_base:
98    where_am_i = subprocess.check_output(
99        ['git', 'rev-parse', '--abbrev-ref', 'HEAD']).decode().strip()
100    # checkout the diff base (="old")
101    subprocess.check_call(['git', 'checkout', args.diff_base])
102    try:
103        old = _run()
104    finally:
105        # restore the original revision (="cur")
106        subprocess.check_call(['git', 'checkout', where_am_i])
107
108text = ''
109if old is None:
110    print(cur)
111    for key, value in sorted(cur.items()):
112        text += '{}: {}\n'.format(key, value)
113else:
114    print(cur, old)
115    call_diff_size = 0
116    channel_diff_size = 0
117    for scenario in _SCENARIOS.keys():
118        for key, value in sorted(_INTERESTING.items()):
119            key = scenario + ': ' + key
120            if key in cur:
121                if key not in old:
122                    text += '{}: {}\n'.format(key, cur[key])
123                else:
124                    text += '{}: {} -> {}\n'.format(key, old[key], cur[key])
125                    if 'call' in key:
126                        call_diff_size += cur[key] - old[key]
127                    else:
128                        channel_diff_size += cur[key] - old[key]
129
130    print("CALL_DIFF_SIZE: %f" % call_diff_size)
131    print("CHANNEL_DIFF_SIZE: %f" % channel_diff_size)
132    check_on_pr.label_increase_decrease_on_pr('per-call-memory', call_diff_size,
133                                              64)
134    check_on_pr.label_increase_decrease_on_pr('per-channel-memory',
135                                              channel_diff_size, 1000)
136    #TODO(chennancy)Change significant value when minstack also runs for channel
137
138print(text)
139check_on_pr.check_on_pr('Memory Difference', '```\n%s\n```' % text)
140