1# Copyright 2018 The gRPC Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Patches the compile() to allow enable parallel compilation of C/C++.
15
16build_ext has lots of C/C++ files and normally them one by one.
17Enabling parallel build helps a lot.
18"""
19
20import distutils.ccompiler
21import os
22
23try:
24    BUILD_EXT_COMPILER_JOBS = int(
25        os.environ['GRPC_PYTHON_BUILD_EXT_COMPILER_JOBS'])
26except KeyError:
27    import multiprocessing
28    BUILD_EXT_COMPILER_JOBS = multiprocessing.cpu_count()
29
30
31# monkey-patch for parallel compilation
32def _parallel_compile(self,
33                      sources,
34                      output_dir=None,
35                      macros=None,
36                      include_dirs=None,
37                      debug=0,
38                      extra_preargs=None,
39                      extra_postargs=None,
40                      depends=None):
41    # setup the same way as distutils.ccompiler.CCompiler
42    # https://github.com/python/cpython/blob/31368a4f0e531c19affe2a1becd25fc316bc7501/Lib/distutils/ccompiler.py#L564
43    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
44        output_dir, macros, include_dirs, sources, depends, extra_postargs)
45    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
46
47    def _compile_single_file(obj):
48        try:
49            src, ext = build[obj]
50        except KeyError:
51            return
52        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
53
54    # run compilation of individual files in parallel
55    import multiprocessing.pool
56    multiprocessing.pool.ThreadPool(BUILD_EXT_COMPILER_JOBS).map(
57        _compile_single_file, objects)
58    return objects
59
60
61def monkeypatch_compile_maybe():
62    """Monkeypatching is dumb, but the build speed gain is worth it."""
63    if BUILD_EXT_COMPILER_JOBS > 1:
64        distutils.ccompiler.CCompiler.compile = _parallel_compile
65