Merge pull request #48636 from terminalmage/loader-fixes

Properly handle optimized .pyc files in loader in PY3

doc/ref/configuration/master.rst

@@ -1612,6 +1612,25 @@ the master will drop the request and the minion's key will remain accepted.
 
     allow_minion_key_revoke: False
 
+.. conf_master:: optimization_order
+
+``optimization_order``
+----------------------
+
+Default: ``[0, 1, 2]``
+
+In cases where Salt is distributed without .py files, this option determines
+the priority of optimization level(s) Salt's module loader should prefer.
+
+.. note::
+    This option is only supported on Python 3.5+.
+
+.. code-block:: yaml
+
+    optimization_order:
+      - 2
+      - 0
+      - 1
 
 Master Large Scale Tuning Settings
 ==================================

doc/ref/configuration/minion.rst

@@ -1195,6 +1195,25 @@ option to ``False`` to keep Salt from updating the mine with this information.
 
     docker.update_mine: False
 
+.. conf_minion:: optimization_order
+
+``optimization_order``
+----------------------
+
+Default: ``[0, 1, 2]``
+
+In cases where Salt is distributed without .py files, this option determines
+the priority of optimization level(s) Salt's module loader should prefer.
+
+.. note::
+    This option is only supported on Python 3.5+.
+
+.. code-block:: yaml
+
+    optimization_order:
+      - 2
+      - 0
+      - 1
 
 Minion Module Management
 ========================

salt/config/__init__.py

@@ -286,6 +286,9 @@ VALID_OPTS = {
     # The type of hashing algorithm to use when doing file comparisons
     'hash_type': str,
 
+    # Order of preference for optimized .pyc files (PY3 only)
+    'optimization_order': list,
+
     # Refuse to load these modules
     'disable_modules': list,
 
@@ -1218,6 +1221,7 @@ DEFAULT_MINION_OPTS = {
     'gitfs_saltenv': [],
     'gitfs_refspecs': _DFLT_REFSPECS,
     'hash_type': 'sha256',
+    'optimization_order': [0, 1, 2],
     'disable_modules': [],
     'disable_returners': [],
     'whitelist_modules': [],
@@ -1511,6 +1515,7 @@ DEFAULT_MASTER_OPTS = {
     'fileserver_verify_config': True,
     'max_open_files': 100000,
     'hash_type': 'sha256',
+    'optimization_order': [0, 1, 2],
     'conf_file': os.path.join(salt.syspaths.CONFIG_DIR, 'master'),
     'open_mode': False,
     'auto_accept': False,

salt/loader.py

@@ -8,6 +8,7 @@ plugin interfaces used by Salt.
 # Import python libs
 from __future__ import absolute_import
 import os
+import re
 import sys
 import time
 import logging
@@ -79,8 +80,8 @@ if USE_IMPORTLIB:
 else:
     SUFFIXES = imp.get_suffixes()
 
-BIN_PRE_EXT = '' if six.PY2 \
-    else '.cpython-{0}{1}'.format(sys.version_info.major, sys.version_info.minor)
+PY3_PRE_EXT = \
+    re.compile(r'\.cpython-{0}{1}(\.opt-[1-9])?'.format(*sys.version_info[:2]))
 
 # Because on the cloud drivers we do `from salt.cloud.libcloudfuncs import *`
 # which simplifies code readability, it adds some unsupported functions into
@@ -1103,6 +1104,15 @@ class LazyLoader(salt.utils.lazy.LazyDict):
         self.disabled = set(self.opts.get('disable_{0}{1}'.format(
             self.tag, '' if self.tag[-1] == 's' else 's'), []))
 
+        # A map of suffix to description for imp
+        self.suffix_map = {}
+        # A list to determine precedence of extensions
+        # Prefer packages (directories) over modules (single files)!
+        self.suffix_order = ['']
+        for (suffix, mode, kind) in SUFFIXES:
+            self.suffix_map[suffix] = (suffix, mode, kind)
+            self.suffix_order.append(suffix)
+
         self._lock = threading.RLock()
         self._refresh_file_mapping()
 
@@ -1176,14 +1186,6 @@ class LazyLoader(salt.utils.lazy.LazyDict):
         refresh the mapping of the FS on disk
         '''
         # map of suffix to description for imp
-        self.suffix_map = {}
-        suffix_order = ['']  # local list to determine precedence of extensions
-                             # Prefer packages (directories) over modules (single files)!
-
-        for (suffix, mode, kind) in SUFFIXES:
-            self.suffix_map[suffix] = (suffix, mode, kind)
-            suffix_order.append(suffix)
-
         if self.opts.get('cython_enable', True) is True:
             try:
                 global pyximport
@@ -1207,11 +1209,25 @@ class LazyLoader(salt.utils.lazy.LazyDict):
         # The files are added in order of priority, so order *must* be retained.
         self.file_mapping = salt.utils.odict.OrderedDict()
 
+        opt_match = []
+
+        def _replace_pre_ext(obj):
+            '''
+            Hack so we can get the optimization level that we replaced (if
+            any) out of the re.sub call below. We use a list here because
+            it is a persistent data structure that we will be able to
+            access after re.sub is called.
+            '''
+            opt_match.append(obj)
+            return ''
+
         for mod_dir in self.module_dirs:
             try:
                 # Make sure we have a sorted listdir in order to have
                 # expectable override results
-                files = sorted(os.listdir(mod_dir))
+                files = sorted(
+                    x for x in os.listdir(mod_dir) if x != '__pycache__'
+                )
             except OSError:
                 continue  # Next mod_dir
             if six.PY3:
@@ -1223,8 +1239,8 @@ class LazyLoader(salt.utils.lazy.LazyDict):
                 except OSError:
                     pass
                 else:
-                    pycache_files.extend(files)
-                    files = pycache_files
+                    files.extend(pycache_files)
+
             for filename in files:
                 try:
                     dirname, basename = os.path.split(filename)
@@ -1233,7 +1249,30 @@ class LazyLoader(salt.utils.lazy.LazyDict):
                         # log messages omitted for obviousness
                         continue  # Next filename
                     f_noext, ext = os.path.splitext(basename)
-                    f_noext = f_noext.replace(BIN_PRE_EXT, '')
+                    if six.PY3:
+                        f_noext = PY3_PRE_EXT.sub(_replace_pre_ext, f_noext)
+                        try:
+                            opt_level = int(
+                                opt_match.pop().group(1).rsplit('-', 1)[-1]
+                            )
+                        except (AttributeError, IndexError, ValueError):
+                            # No regex match or no optimization level matched
+                            opt_level = 0
+                        try:
+                            opt_index = self.opts['optimization_order'].index(opt_level)
+                        except KeyError:
+                            log.trace(
+                                'Disallowed optimization level %d for module '
+                                'name \'%s\', skipping. Add %d to the '
+                                '\'optimization_order\' config option if you '
+                                'do not want to ignore this optimization '
+                                'level.', opt_level, f_noext, opt_level
+                            )
+                            continue
+                    else:
+                        # Optimization level not reflected in filename on PY2
+                        opt_index = 0
+
                     # make sure it is a suffix we support
                     if ext not in self.suffix_map:
                         continue  # Next filename
@@ -1249,7 +1288,7 @@ class LazyLoader(salt.utils.lazy.LazyDict):
                     if ext == '':
                         # is there something __init__?
                         subfiles = os.listdir(fpath)
-                        for suffix in suffix_order:
+                        for suffix in self.suffix_order:
                             if '' == suffix:
                                 continue  # Next suffix (__init__ must have a suffix)
                             init_file = '__init__{0}'.format(suffix)
@@ -1258,17 +1297,29 @@ class LazyLoader(salt.utils.lazy.LazyDict):
                         else:
                             continue  # Next filename
 
-                    if f_noext in self.file_mapping:
+                    try:
                         curr_ext = self.file_mapping[f_noext][1]
-                        #log.debug("****** curr_ext={0} ext={1} suffix_order={2}".format(curr_ext, ext, suffix_order))
+                        curr_opt_index = self.file_mapping[f_noext][2]
+                    except KeyError:
+                        pass
+                    else:
                         if '' in (curr_ext, ext) and curr_ext != ext:
                             log.error(
                                 'Module/package collision: \'%s\' and \'%s\'',
                                 fpath,
                                 self.file_mapping[f_noext][0]
                             )
-                        if not curr_ext or suffix_order.index(ext) >= suffix_order.index(curr_ext):
-                            continue  # Next filename
+
+                        if six.PY3 and ext == '.pyc' and curr_ext == '.pyc':
+                            # Check the optimization level
+                            if opt_index >= curr_opt_index:
+                                # Module name match, but a higher-priority
+                                # optimization level was already matched, skipping.
+                                continue
+                        elif not curr_ext or self.suffix_order.index(ext) >= self.suffix_order.index(curr_ext):
+                            # Match found but a higher-priorty match already
+                            # exists, so skip this.
+                            continue
 
                     if six.PY3 and not dirname and ext == '.pyc':
                         # On Python 3, we should only load .pyc files from the
@@ -1277,13 +1328,13 @@ class LazyLoader(salt.utils.lazy.LazyDict):
                         continue
 
                     # Made it this far - add it
-                    self.file_mapping[f_noext] = (fpath, ext)
+                    self.file_mapping[f_noext] = (fpath, ext, opt_index)
 
                 except OSError:
                     continue
         for smod in self.static_modules:
             f_noext = smod.split('.')[-1]
-            self.file_mapping[f_noext] = (smod, '.o')
+            self.file_mapping[f_noext] = (smod, '.o', 0)
 
     def clear(self):
         '''
@@ -1352,7 +1403,7 @@ class LazyLoader(salt.utils.lazy.LazyDict):
 
     def _load_module(self, name):
         mod = None
-        fpath, suffix = self.file_mapping[name]
+        fpath, suffix = self.file_mapping[name][:2]
         self.loaded_files.add(name)
         fpath_dirname = os.path.dirname(fpath)
         try:

tests/integration/minion/test_pillar.py

@@ -37,6 +37,7 @@ GPG_SLS = os.path.join(PILLAR_BASE, 'gpg.sls')
 DEFAULT_OPTS = {
     'cachedir': os.path.join(TMP, 'rootdir', 'cache'),
     'config_dir': TMP_CONF_DIR,
+    'optimization_order': [0, 1, 2],
     'extension_modules': os.path.join(TMP,
                                       'test-decrypt-pillar',
                                       'extmods'),

tests/support/gitfs.py

@@ -330,11 +330,14 @@ class GitPillarTestBase(GitTestBase, LoaderModuleMockMixin):
         '''
         cachedir = tempfile.mkdtemp(dir=TMP)
         self.addCleanup(shutil.rmtree, cachedir, ignore_errors=True)
-        ext_pillar_opts = yaml.safe_load(
-            ext_pillar_conf.format(
-                cachedir=cachedir,
-                extmods=os.path.join(cachedir, 'extmods'),
-                **self.ext_opts
+        ext_pillar_opts = {'optimization_order': [0, 1, 2]}
+        ext_pillar_opts.update(
+            yaml.safe_load(
+                ext_pillar_conf.format(
+                    cachedir=cachedir,
+                    extmods=os.path.join(cachedir, 'extmods'),
+                    **self.ext_opts
+                )
             )
         )
         with patch.dict(git_pillar.__opts__, ext_pillar_opts):

tests/unit/loader/test_loader.py

@@ -8,9 +8,11 @@
 
 # Import Python libs
 from __future__ import absolute_import
+import compileall
 import inspect
 import logging
 import tempfile
+import textwrap
 import shutil
 import os
 import collections
@@ -31,7 +33,15 @@ import salt.ext.six as six
 from salt.ext.six.moves import range
 # pylint: enable=no-name-in-module,redefined-builtin
 
-from salt.loader import LazyLoader, _module_dirs, grains, utils, proxy, minion_mods
+from salt.loader import (
+    LazyLoader,
+    USE_IMPORTLIB,
+    _module_dirs,
+    grains,
+    utils,
+    proxy,
+    minion_mods
+)
 
 log = logging.getLogger(__name__)
 
@@ -916,3 +926,131 @@ class LazyLoaderDeepSubmodReloadingTest(TestCase):
                 self.update_lib(lib)
                 self.loader.clear()
                 self._verify_libs()
+
+
+class LazyLoaderOptimizationOrderTest(TestCase):
+    '''
+    Test the optimization order priority in the loader (PY3)
+    '''
+    module_name = 'lazyloadertest'
+    module_content = textwrap.dedent('''\
+        # -*- coding: utf-8 -*-
+        from __future__ import absolute_import
+
+        def test():
+            return True
+        ''')
+
+    @classmethod
+    def setUpClass(cls):
+        cls.opts = salt.config.minion_config(None)
+        cls.opts['grains'] = grains(cls.opts)
+
+    def setUp(self):
+        # Setup the module
+        self.module_dir = tempfile.mkdtemp(dir=TMP)
+        self.module_file = os.path.join(self.module_dir,
+                                        '{0}.py'.format(self.module_name))
+
+    def _get_loader(self, order=None):
+        opts = copy.deepcopy(self.opts)
+        if order is not None:
+            opts['optimization_order'] = order
+        # Return a loader
+        return LazyLoader([self.module_dir], opts, tag='module')
+
+    def _get_module_filename(self):
+        # The act of referencing the loader entry forces the module to be
+        # loaded by the LazyDict.
+        mod_fullname = self.loader[next(iter(self.loader))].__module__
+        return sys.modules[mod_fullname].__file__
+
+    def _expected(self, optimize=0):
+        if six.PY3:
+            return 'lazyloadertest.cpython-{0}{1}{2}.pyc'.format(
+                sys.version_info[0],
+                sys.version_info[1],
+                '' if not optimize else '.opt-{0}'.format(optimize)
+            )
+        else:
+            return 'lazyloadertest.pyc'
+
+    def _write_module_file(self):
+        with salt.utils.fopen(self.module_file, 'w') as fh:
+            fh.write(self.module_content)
+            fh.flush()
+            os.fsync(fh.fileno())
+
+    def _byte_compile(self):
+        if USE_IMPORTLIB:
+            # Skip this check as "optimize" is unique to PY3's compileall
+            # module, and this will be a false error when Pylint is run on
+            # Python 2.
+            # pylint: disable=unexpected-keyword-arg
+            compileall.compile_file(self.module_file, quiet=1, optimize=0)
+            compileall.compile_file(self.module_file, quiet=1, optimize=1)
+            compileall.compile_file(self.module_file, quiet=1, optimize=2)
+            # pylint: enable=unexpected-keyword-arg
+        else:
+            compileall.compile_file(self.module_file, quiet=1)
+
+    def _test_optimization_order(self, order):
+        self._write_module_file()
+        self._byte_compile()
+
+        # Clean up the original file so that we can be assured we're only
+        # loading the byte-compiled files(s).
+        os.remove(self.module_file)
+
+        self.loader = self._get_loader(order)
+        filename = self._get_module_filename()
+        basename = os.path.basename(filename)
+        assert basename == self._expected(order[0]), basename
+
+        if not USE_IMPORTLIB:
+            # We are only testing multiple optimization levels on Python 3.5+
+            return
+
+        # Remove the file and make a new loader. We should now load the
+        # byte-compiled file with an optimization level matching the 2nd
+        # element of the order list.
+        os.remove(filename)
+        self.loader = self._get_loader(order)
+        filename = self._get_module_filename()
+        basename = os.path.basename(filename)
+        assert basename == self._expected(order[1]), basename
+
+        # Remove the file and make a new loader. We should now load the
+        # byte-compiled file with an optimization level matching the 3rd
+        # element of the order list.
+        os.remove(filename)
+        self.loader = self._get_loader(order)
+        filename = self._get_module_filename()
+        basename = os.path.basename(filename)
+        assert basename == self._expected(order[2]), basename
+
+    def test_optimization_order(self):
+        '''
+        Test the optimization_order config param
+        '''
+        self._test_optimization_order([0, 1, 2])
+        self._test_optimization_order([0, 2, 1])
+        if USE_IMPORTLIB:
+            # optimization_order only supported on Python 3.5+, earlier
+            # releases only support unoptimized .pyc files.
+            self._test_optimization_order([1, 2, 0])
+            self._test_optimization_order([1, 0, 2])
+            self._test_optimization_order([2, 0, 1])
+            self._test_optimization_order([2, 1, 0])
+
+    def test_load_source_file(self):
+        '''
+        Make sure that .py files are preferred over .pyc files
+        '''
+        self._write_module_file()
+        self._byte_compile()
+        self.loader = self._get_loader()
+        filename = self._get_module_filename()
+        basename = os.path.basename(filename)
+        expected = 'lazyloadertest.py' if six.PY3 else 'lazyloadertest.pyc'
+        assert basename == expected, basename

tests/unit/output/test_highstate_out.py

@@ -27,6 +27,7 @@ class JsonTestCase(TestCase, LoaderModuleMockMixin):
             highstate: {
                 '__opts__': {
                     'extension_modules': '',
+                    'optimization_order': [0, 1, 2],
                     'color': False,
                 }
             }

tests/unit/pillar/test_git.py

@@ -60,6 +60,7 @@ class GitPillarTestCase(TestCase, AdaptedConfigurationTestCaseMixin, LoaderModul
                     'file_roots': {},
                     'state_top': 'top.sls',
                     'extension_modules': '',
+                    'optimization_order': [0, 1, 2],
                     'renderer': 'yaml_jinja',
                     'renderer_blacklist': [],
                     'renderer_whitelist': [],

tests/unit/returners/test_smtp_return.py

@@ -31,7 +31,25 @@ class SMTPReturnerTestCase(TestCase, LoaderModuleMockMixin):
     Test SMTP returner
     '''
     def setup_loader_modules(self):
-        return {smtp: {}}
+        return {
+            smtp: {
+                '__opts__': {
+                    'extension_modules': '',
+                    'optimization_order': [0, 1, 2],
+                    'renderer': 'jinja|yaml',
+                    'renderer_blacklist': [],
+                    'renderer_whitelist': [],
+                    'file_roots': {},
+                    'pillar_roots': {},
+                    'cachedir': '/',
+                    'master_uri': 'tcp://127.0.0.1:4505',
+                    'pki_dir': '/',
+                    'keysize': 2048,
+                    'id': 'test',
+                    '__role': 'minion',
+                }
+            }
+        }
 
     def _test_returner(self, mocked_smtplib):  # pylint: disable=unused-argument
         '''
@@ -60,25 +78,11 @@ class SMTPReturnerTestCase(TestCase, LoaderModuleMockMixin):
 
     if HAS_GNUPG:
         def test_returner(self):
-            with patch.dict(smtp.__opts__, {'extension_modules': '',
-                                            'renderer': 'jinja|yaml',
-                                            'renderer_blacklist': [],
-                                            'renderer_whitelist': [],
-                                            'file_roots': {},
-                                            'pillar_roots': {},
-                                            'cachedir': '/'}), \
-                    patch('salt.returners.smtp_return.gnupg'), \
+            with patch('salt.returners.smtp_return.gnupg'), \
                     patch('salt.returners.smtp_return.smtplib.SMTP') as mocked_smtplib:
                 self._test_returner(mocked_smtplib)
 
     else:
         def test_returner(self):
-            with patch.dict(smtp.__opts__, {'extension_modules': '',
-                                            'renderer': 'jinja|yaml',
-                                            'renderer_blacklist': [],
-                                            'renderer_whitelist': [],
-                                            'file_roots': {},
-                                            'pillar_roots': {},
-                                            'cachedir': '/'}), \
-                    patch('salt.returners.smtp_return.smtplib.SMTP') as mocked_smtplib:
+            with patch('salt.returners.smtp_return.smtplib.SMTP') as mocked_smtplib:
                 self._test_returner(mocked_smtplib)

tests/unit/runners/test_winrepo.py

@@ -86,6 +86,7 @@ class WinrepoTest(TestCase, LoaderModuleMockMixin):
             winrepo: {
                 '__opts__': {
                     'winrepo_cachefile': 'winrepo.p',
+                    'optimization_order': [0, 1, 2],
                     'renderer': 'yaml',
                     'renderer_blacklist': [],
                     'renderer_whitelist': [],

tests/unit/test_map_conf.py

@@ -85,6 +85,7 @@ class MapConfTest(TestCase):
                 patch('salt.cloud.Map.read', MagicMock(return_value=EXAMPLE_MAP)):
             self.maxDiff = None
             opts = {'extension_modules': '/var/cache/salt/master/extmods',
+                    'optimization_order': [0, 1, 2],
                     'providers': EXAMPLE_PROVIDERS, 'profiles': EXAMPLE_PROFILES}
             cloud_map = salt.cloud.Map(opts)
             merged_profile = {

tests/unit/test_pillar.py

@@ -33,6 +33,7 @@ class PillarTestCase(TestCase):
     def test_pillarenv_from_saltenv(self):
         with patch('salt.pillar.compile_template') as compile_template:
             opts = {
+                'optimization_order': [0, 1, 2],
                 'renderer': 'json',
                 'renderer_blacklist': [],
                 'renderer_whitelist': [],
@@ -51,6 +52,7 @@ class PillarTestCase(TestCase):
 
     def test_dynamic_pillarenv(self):
         opts = {
+            'optimization_order': [0, 1, 2],
             'renderer': 'json',
             'renderer_blacklist': [],
             'renderer_whitelist': [],
@@ -65,6 +67,7 @@ class PillarTestCase(TestCase):
 
     def test_ignored_dynamic_pillarenv(self):
         opts = {
+            'optimization_order': [0, 1, 2],
             'renderer': 'json',
             'renderer_blacklist': [],
             'renderer_whitelist': [],
@@ -79,6 +82,7 @@ class PillarTestCase(TestCase):
     def test_malformed_pillar_sls(self):
         with patch('salt.pillar.compile_template') as compile_template:
             opts = {
+                'optimization_order': [0, 1, 2],
                 'renderer': 'json',
                 'renderer_blacklist': [],
                 'renderer_whitelist': [],
@@ -188,6 +192,7 @@ class PillarTestCase(TestCase):
 
     def test_includes_override_sls(self):
         opts = {
+            'optimization_order': [0, 1, 2],
             'renderer': 'json',
             'renderer_blacklist': [],
             'renderer_whitelist': [],
@@ -250,6 +255,7 @@ class PillarTestCase(TestCase):
         with patch('salt.pillar.salt.fileclient.get_file_client', autospec=True) as get_file_client, \
                 patch('salt.pillar.salt.minion.Matcher') as Matcher:  # autospec=True disabled due to py3 mock bug
             opts = {
+                'optimization_order': [0, 1, 2],
                 'renderer': 'yaml',
                 'renderer_blacklist': [],
                 'renderer_whitelist': [],