[WIP] relocate.py: parallelize test replacement logic (#19690)

* sbang pushed back to callers;
star moved to util.lang

* updated unit test

* sbang test moved; local tests pass

Co-authored-by: Nathan Hanford <hanford1@llnl.gov>

10 files changed, 316 insertions, 210 deletions

diff --git a/lib/spack/llnl/util/lang.py b/lib/spack/llnl/util/lang.py
index 88058d4bd1..8d7d5e0767 100644
--- a/lib/spack/llnl/util/lang.py
+++ b/lib/spack/llnl/util/lang.py
@@ -673,6 +673,13 @@ def uniq(sequence):
     return uniq_list
 
 
+def star(func):
+    """Unpacks arguments for use with Multiprocessing mapping functions"""
+    def _wrapper(args):
+        return func(*args)
+    return _wrapper
+
+
 class Devnull(object):
     """Null stream with less overhead than ``os.devnull``.
 
diff --git a/lib/spack/spack/binary_distribution.py b/lib/spack/spack/binary_distribution.py
index d643bde7b3..a5fb94a770 100644
--- a/lib/spack/spack/binary_distribution.py
+++ b/lib/spack/spack/binary_distribution.py
@@ -12,6 +12,7 @@ import shutil
 import tempfile
 import hashlib
 import glob
+from ordereddict_backport import OrderedDict
 
 from contextlib import closing
 import ruamel.yaml as yaml
@@ -1105,11 +1106,26 @@ def relocate_package(spec, allow_root):
     new_deps = spack.build_environment.get_rpath_deps(spec.package)
     for d in new_deps:
         hash_to_prefix[d.format('{hash}')] = str(d.prefix)
-    prefix_to_prefix = dict()
+    # Spurious replacements (e.g. sbang) will cause issues with binaries
+    # For example, the new sbang can be longer than the old one.
+    # Hence 2 dictionaries are maintained here.
+    prefix_to_prefix_text = OrderedDict({})
+    prefix_to_prefix_bin = OrderedDict({})
+    prefix_to_prefix_text[old_prefix] = new_prefix
+    prefix_to_prefix_bin[old_prefix] = new_prefix
+    prefix_to_prefix_text[old_layout_root] = new_layout_root
+    prefix_to_prefix_bin[old_layout_root] = new_layout_root
     for orig_prefix, hash in prefix_to_hash.items():
-        prefix_to_prefix[orig_prefix] = hash_to_prefix.get(hash, None)
-    prefix_to_prefix[old_prefix] = new_prefix
-    prefix_to_prefix[old_layout_root] = new_layout_root
+        prefix_to_prefix_text[orig_prefix] = hash_to_prefix.get(hash, None)
+        prefix_to_prefix_bin[orig_prefix] = hash_to_prefix.get(hash, None)
+    # This is vestigial code for the *old* location of sbang. Previously,
+    # sbang was a bash script, and it lived in the spack prefix. It is
+    # now a POSIX script that lives in the install prefix. Old packages
+    # will have the old sbang location in their shebangs.
+    import spack.hooks.sbang as sbang
+    orig_sbang = '#!/bin/bash {0}/bin/sbang'.format(old_spack_prefix)
+    new_sbang = sbang.sbang_shebang_line()
+    prefix_to_prefix_text[orig_sbang] = new_sbang
 
     tty.debug("Relocating package from",
               "%s to %s." % (old_layout_root, new_layout_root))
@@ -1137,15 +1153,14 @@ def relocate_package(spec, allow_root):
             relocate.relocate_macho_binaries(files_to_relocate,
                                              old_layout_root,
                                              new_layout_root,
-                                             prefix_to_prefix, rel,
+                                             prefix_to_prefix_bin, rel,
                                              old_prefix,
                                              new_prefix)
-
         if 'elf' in platform.binary_formats:
             relocate.relocate_elf_binaries(files_to_relocate,
                                            old_layout_root,
                                            new_layout_root,
-                                           prefix_to_prefix, rel,
+                                           prefix_to_prefix_bin, rel,
                                            old_prefix,
                                            new_prefix)
             # Relocate links to the new install prefix
@@ -1156,12 +1171,7 @@ def relocate_package(spec, allow_root):
 
         # For all buildcaches
         # relocate the install prefixes in text files including dependencies
-        relocate.relocate_text(text_names,
-                               old_layout_root, new_layout_root,
-                               old_prefix, new_prefix,
-                               old_spack_prefix,
-                               new_spack_prefix,
-                               prefix_to_prefix)
+        relocate.relocate_text(text_names, prefix_to_prefix_text)
 
         paths_to_relocate = [old_prefix, old_layout_root]
         paths_to_relocate.extend(prefix_to_hash.keys())
@@ -1171,22 +1181,13 @@ def relocate_package(spec, allow_root):
             map(lambda filename: os.path.join(workdir, filename),
                 buildinfo['relocate_binaries'])))
         # relocate the install prefixes in binary files including dependencies
-        relocate.relocate_text_bin(files_to_relocate,
-                                   old_prefix, new_prefix,
-                                   old_spack_prefix,
-                                   new_spack_prefix,
-                                   prefix_to_prefix)
-
-# If we are installing back to the same location
-# relocate the sbang location if the spack directory changed
+        relocate.relocate_text_bin(files_to_relocate, prefix_to_prefix_bin)
+
+    # If we are installing back to the same location
+    # relocate the sbang location if the spack directory changed
     else:
         if old_spack_prefix != new_spack_prefix:
-            relocate.relocate_text(text_names,
-                                   old_layout_root, new_layout_root,
-                                   old_prefix, new_prefix,
-                                   old_spack_prefix,
-                                   new_spack_prefix,
-                                   prefix_to_prefix)
+            relocate.relocate_text(text_names, prefix_to_prefix_text)
 
 
 def extract_tarball(spec, filename, allow_root=False, unsigned=False,
diff --git a/lib/spack/spack/filesystem_view.py b/lib/spack/spack/filesystem_view.py
index 0f21e4d975..b6501064b9 100644
--- a/lib/spack/spack/filesystem_view.py
+++ b/lib/spack/spack/filesystem_view.py
@@ -8,6 +8,7 @@ import os
 import re
 import shutil
 import sys
+from ordereddict_backport import OrderedDict
 
 from llnl.util.link_tree import LinkTree, MergeConflictError
 from llnl.util import tty
@@ -65,32 +66,35 @@ def view_copy(src, dst, view, spec=None):
         # Not metadata, we have to relocate it
 
         # Get information on where to relocate from/to
-        prefix_to_projection = dict(
-            (dep.prefix, view.get_projection_for_spec(dep))
-            for dep in spec.traverse()
-        )
+
+        # This is vestigial code for the *old* location of sbang. Previously,
+        # sbang was a bash script, and it lived in the spack prefix. It is
+        # now a POSIX script that lives in the install prefix. Old packages
+        # will have the old sbang location in their shebangs.
+        # TODO: Not sure which one to use...
+        import spack.hooks.sbang as sbang
+        orig_sbang = '#!/bin/bash {0}/bin/sbang'.format(spack.paths.spack_root)
+        new_sbang = sbang.sbang_shebang_line()
+
+        prefix_to_projection = OrderedDict({
+            spec.prefix: view.get_projection_for_spec(spec),
+            spack.paths.spack_root: view._root})
+
+        for dep in spec.traverse():
+            prefix_to_projection[dep.prefix] = \
+                view.get_projection_for_spec(dep)
 
         if spack.relocate.is_binary(dst):
-            # relocate binaries
             spack.relocate.relocate_text_bin(
                 binaries=[dst],
-                orig_install_prefix=spec.prefix,
-                new_install_prefix=view.get_projection_for_spec(spec),
-                orig_spack=spack.paths.spack_root,
-                new_spack=view._root,
-                new_prefixes=prefix_to_projection
+                prefixes=prefix_to_projection
             )
         else:
-            # relocate text
+            prefix_to_projection[spack.store.layout.root] = view._root
+            prefix_to_projection[orig_sbang] = new_sbang
             spack.relocate.relocate_text(
                 files=[dst],
-                orig_layout_root=spack.store.layout.root,
-                new_layout_root=view._root,
-                orig_install_prefix=spec.prefix,
-                new_install_prefix=view.get_projection_for_spec(spec),
-                orig_spack=spack.paths.spack_root,
-                new_spack=view._root,
-                new_prefixes=prefix_to_projection
+                prefixes=prefix_to_projection
             )
 
 
diff --git a/lib/spack/spack/relocate.py b/lib/spack/spack/relocate.py
index 2685b0c6ab..e1726b060e 100644
--- a/lib/spack/spack/relocate.py
+++ b/lib/spack/spack/relocate.py
@@ -6,6 +6,8 @@ import os
 import platform
 import re
 import shutil
+import multiprocessing.pool
+from ordereddict_backport import OrderedDict
 
 import llnl.util.lang
 import llnl.util.tty as tty
@@ -449,36 +451,26 @@ def needs_text_relocation(m_type, m_subtype):
     return m_type == 'text'
 
 
-def _replace_prefix_text(filename, old_dir, new_dir):
+def _replace_prefix_text(filename, compiled_prefixes):
     """Replace all the occurrences of the old install prefix with a
     new install prefix in text files that are utf-8 encoded.
 
     Args:
         filename (str): target text file (utf-8 encoded)
-        old_dir (str): directory to be searched in the file
-        new_dir (str): substitute for the old directory
+        compiled_prefixes (OrderedDict): OrderedDictionary where the keys are
+        precompiled regex of the old prefixes and the values are the new
+        prefixes (uft-8 encoded)
     """
-    # TODO: cache regexes globally to speedup computation
     with open(filename, 'rb+') as f:
         data = f.read()
         f.seek(0)
-        # Replace old_dir with new_dir if it appears at the beginning of a path
-        # Negative lookbehind for a character legal in a path
-        # Then a match group for any characters legal in a compiler flag
-        # Then old_dir
-        # Then characters legal in a path
-        # Ensures we only match the old_dir if it's precedeed by a flag or by
-        # characters not legal in a path, but not if it's preceeded by other
-        # components of a path.
-        old_bytes = old_dir.encode('utf-8')
-        pat = b'(?<![\\w\\-_/])([\\w\\-_]*?)%s([\\w\\-_/]*)' % old_bytes
-        repl = b'\\1%s\\2' % new_dir.encode('utf-8')
-        ndata = re.sub(pat, repl, data)
-        f.write(ndata)
+        for orig_prefix_rexp, new_bytes in compiled_prefixes.items():
+            data = orig_prefix_rexp.sub(new_bytes, data)
+        f.write(data)
         f.truncate()
 
 
-def _replace_prefix_bin(filename, old_dir, new_dir):
+def _replace_prefix_bin(filename, byte_prefixes):
     """Replace all the occurrences of the old install prefix with a
     new install prefix in binary files.
 
@@ -487,33 +479,34 @@ def _replace_prefix_bin(filename, old_dir, new_dir):
 
     Args:
         filename (str): target binary file
-        old_dir (str): directory to be searched in the file
-        new_dir (str): substitute for the old directory
-    """
-    def replace(match):
-        occurances = match.group().count(old_dir.encode('utf-8'))
-        olen = len(old_dir.encode('utf-8'))
-        nlen = len(new_dir.encode('utf-8'))
-        padding = (olen - nlen) * occurances
-        if padding < 0:
-            return data
-        return match.group().replace(
-            old_dir.encode('utf-8'),
-            os.sep.encode('utf-8') * padding + new_dir.encode('utf-8')
-        )
+        byte_prefixes (OrderedDict): OrderedDictionary where the keys are
+        precompiled regex of the old prefixes and the values are the new
+        prefixes (uft-8 encoded)
+    """
 
     with open(filename, 'rb+') as f:
         data = f.read()
         f.seek(0)
-        original_data_len = len(data)
-        pat = re.compile(old_dir.encode('utf-8'))
-        if not pat.search(data):
-            return
-        ndata = pat.sub(replace, data)
-        if not len(ndata) == original_data_len:
-            raise BinaryStringReplacementError(
-                filename, original_data_len, len(ndata))
-        f.write(ndata)
+        for orig_bytes, new_bytes in byte_prefixes.items():
+            original_data_len = len(data)
+            # Skip this hassle if not found
+            if orig_bytes not in data:
+                continue
+            # We only care about this problem if we are about to replace
+            length_compatible = len(new_bytes) <= len(orig_bytes)
+            if not length_compatible:
+                raise BinaryTextReplaceError(orig_bytes, new_bytes)
+            pad_length = len(orig_bytes) - len(new_bytes)
+            padding = os.sep * pad_length
+            padding = padding.encode('utf-8')
+            data = data.replace(orig_bytes, new_bytes + padding)
+            # Really needs to be the same length
+            if not len(data) == original_data_len:
+                print('Length of pad:', pad_length, 'should be', len(padding))
+                print(new_bytes, 'was to replace', orig_bytes)
+                raise BinaryStringReplacementError(
+                    filename, original_data_len, len(data))
+        f.write(data)
         f.truncate()
 
 
@@ -786,86 +779,88 @@ def relocate_links(links, orig_layout_root,
             tty.warn(msg.format(link_target, abs_link, new_install_prefix))
 
 
-def relocate_text(
-        files, orig_layout_root, new_layout_root, orig_install_prefix,
-        new_install_prefix, orig_spack, new_spack, new_prefixes
-):
-    """Relocate text file from the original ``install_tree`` to the new one.
+def relocate_text(files, prefixes, concurrency=32):
+    """Relocate text file from the original installation prefix to the
+     new prefix.
 
-    This also handles relocating Spack's sbang scripts to point at the
-    new install tree.
-
-    Args:
-        files (list): text files to be relocated
-        orig_layout_root (str): original layout root
-        new_layout_root (str): new layout root
-        orig_install_prefix (str): install prefix of the original installation
-        new_install_prefix (str): install prefix where we want to relocate
-        orig_spack (str): path to the original Spack
-        new_spack (str): path to the new Spack
-        new_prefixes (dict): dictionary that maps the original prefixes to
-            where they should be relocated
+     Relocation also affects the the path in Spack's sbang script.
 
+     Args:
+         files (list): Text files to be relocated
+         prefixes (OrderedDict): String prefixes which need to be changed
+         concurrency (int): Preferred degree of parallelism
     """
-    # TODO: reduce the number of arguments (8 seems too much)
 
-    # This is vestigial code for the *old* location of sbang. Previously,
-    # sbang was a bash script, and it lived in the spack prefix. It is
-    # now a POSIX script that lives in the install prefix. Old packages
-    # will have the old sbang location in their shebangs.
-    import spack.hooks.sbang as sbang
-    orig_sbang = '#!/bin/bash {0}/bin/sbang'.format(orig_spack)
-    new_sbang = sbang.sbang_shebang_line()
+    # This now needs to be handled by the caller in all cases
+    # orig_sbang = '#!/bin/bash {0}/bin/sbang'.format(orig_spack)
+    # new_sbang = '#!/bin/bash {0}/bin/sbang'.format(new_spack)
+
+    compiled_prefixes = OrderedDict({})
+
+    for orig_prefix, new_prefix in prefixes.items():
+        if orig_prefix != new_prefix:
+            orig_bytes = orig_prefix.encode('utf-8')
+            orig_prefix_rexp = re.compile(
+                b'(?<![\\w\\-_/])([\\w\\-_]*?)%s([\\w\\-_/]*)' % orig_bytes)
+            new_bytes = b'\\1%s\\2' % new_prefix.encode('utf-8')
+            compiled_prefixes[orig_prefix_rexp] = new_bytes
 
     # Do relocations on text that refers to the install tree
+    # multiprocesing.ThreadPool.map requires single argument
+
+    args = []
     for filename in files:
-        _replace_prefix_text(filename, orig_install_prefix, new_install_prefix)
-        for orig_dep_prefix, new_dep_prefix in new_prefixes.items():
-            _replace_prefix_text(filename, orig_dep_prefix, new_dep_prefix)
-        _replace_prefix_text(filename, orig_layout_root, new_layout_root)
+        args.append((filename, compiled_prefixes))
 
-        # Point old packages at the new sbang location. Packages that
-        # already use the new sbang location will already have been
-        # handled by the prior call to _replace_prefix_text
-        _replace_prefix_text(filename, orig_sbang, new_sbang)
+    tp = multiprocessing.pool.ThreadPool(processes=concurrency)
+    try:
+        tp.map(llnl.util.lang.star(_replace_prefix_text), args)
+    finally:
+        tp.terminate()
+        tp.join()
 
 
-def relocate_text_bin(
-        binaries, orig_install_prefix, new_install_prefix,
-        orig_spack, new_spack, new_prefixes
-):
+def relocate_text_bin(binaries, prefixes, concurrency=32):
     """Replace null terminated path strings hard coded into binaries.
 
     The new install prefix must be shorter than the original one.
 
     Args:
         binaries (list): binaries to be relocated
-        orig_install_prefix (str): install prefix of the original installation
-        new_install_prefix (str): install prefix where we want to relocate
-        orig_spack (str): path to the original Spack
-        new_spack (str): path to the new Spack
-        new_prefixes (dict): dictionary that maps the original prefixes to
-            where they should be relocated
+        prefixes (OrderedDict): String prefixes which need to be changed.
+        concurrency (int): Desired degree of parallelism.
 
     Raises:
-      BinaryTextReplaceError: when the new path in longer than the old path
+      BinaryTextReplaceError: when the new path is longer than the old path
     """
-    # Raise if the new install prefix is longer than the
-    # original one, since it means we can't change the original
-    # binary to relocate it
-    new_prefix_is_shorter = len(new_install_prefix) <= len(orig_install_prefix)
-    if not new_prefix_is_shorter and len(binaries) > 0:
-        raise BinaryTextReplaceError(orig_install_prefix, new_install_prefix)
+    byte_prefixes = OrderedDict({})
+
+    for orig_prefix, new_prefix in prefixes.items():
+        if orig_prefix != new_prefix:
+            if isinstance(orig_prefix, bytes):
+                orig_bytes = orig_prefix
+            else:
+                orig_bytes = orig_prefix.encode('utf-8')
+            if isinstance(new_prefix, bytes):
+                new_bytes = new_prefix
+            else:
+                new_bytes = new_prefix.encode('utf-8')
+            byte_prefixes[orig_bytes] = new_bytes
+
+    # Do relocations on text in binaries that refers to the install tree
+    # multiprocesing.ThreadPool.map requires single argument
+    args = []
 
     for binary in binaries:
-        for old_dep_prefix, new_dep_prefix in new_prefixes.items():
-            if len(new_dep_prefix) <= len(old_dep_prefix):
-                _replace_prefix_bin(binary, old_dep_prefix, new_dep_prefix)
-        _replace_prefix_bin(binary, orig_install_prefix, new_install_prefix)
-
-    # Note: Replacement of spack directory should not be done. This causes
-    # an incorrect replacement path in the case where the install root is a
-    # subdirectory of the spack directory.
+        args.append((binary, byte_prefixes))
+
+    tp = multiprocessing.pool.ThreadPool(processes=concurrency)
+
+    try:
+        tp.map(llnl.util.lang.star(_replace_prefix_bin), args)
+    finally:
+        tp.terminate()
+        tp.join()
 
 
 def is_relocatable(spec):
diff --git a/lib/spack/spack/test/bindist.py b/lib/spack/spack/test/bindist.py
index 974c50c260..1d07fc8483 100644
--- a/lib/spack/spack/test/bindist.py
+++ b/lib/spack/spack/test/bindist.py
@@ -19,6 +19,7 @@ import spack.cmd.buildcache as buildcache
 import spack.cmd.install as install
 import spack.cmd.uninstall as uninstall
 import spack.cmd.mirror as mirror
+import spack.hooks.sbang as sbang
 from spack.main import SpackCommand
 import spack.mirror
 import spack.util.gpg
@@ -80,6 +81,15 @@ def mirror_directory_rel(session_mirror_rel):
     yield(session_mirror_rel)
 
 
+@pytest.fixture(scope='function')
+def function_mirror(tmpdir):
+    mirror_dir = str(tmpdir.join('mirror'))
+    mirror_cmd('add', '--scope', 'site', 'test-mirror-func',
+               'file://%s' % mirror_dir)
+    yield mirror_dir
+    mirror_cmd('rm', '--scope=site', 'test-mirror-func')
+
+
 @pytest.fixture(scope='session')
 def config_directory(tmpdir_factory):
     tmpdir = tmpdir_factory.mktemp('test_configs')
@@ -671,3 +681,76 @@ def test_generate_indices_exception(monkeypatch, capfd):
     err = capfd.readouterr()[1]
     expect = 'Encountered problem listing packages at {0}'.format(test_url)
     assert expect in err
+
+
+@pytest.mark.usefixtures('mock_fetch')
+def test_update_sbang(tmpdir, install_mockery, function_mirror):
+    """
+    Test the creation and installation of buildcaches with default rpaths
+    into the non-default directory layout scheme, triggering an update of the
+    sbang.
+    """
+
+    # Save the original store and layout before we touch ANYTHING.
+    real_store = spack.store.store
+    real_layout = spack.store.layout
+
+    # Concretize a package with some old-fashioned sbang lines.
+    sspec = Spec('old-sbang')
+    sspec.concretize()
+
+    # Need a fake mirror with *function* scope.
+    mirror_dir = function_mirror
+
+    # Assumes all commands will concretize sspec the same way.
+    install_cmd('--no-cache', sspec.name)
+
+    # Create a buildcache with the installed spec.
+    buildcache_cmd('create', '-u', '-a', '-d', mirror_dir,
+                   '/%s' % sspec.dag_hash())
+
+    # Need to force an update of the buildcache index
+    buildcache_cmd('update-index', '-d', 'file://%s' % mirror_dir)
+
+    # Uninstall the original package.
+    uninstall_cmd('-y', '/%s' % sspec.dag_hash())
+
+    try:
+        # New install tree locations...
+        # Too fine-grained to do be done in a fixture
+        spack.store.store = spack.store.Store(str(tmpdir.join('newtree')))
+        spack.store.layout = YamlDirectoryLayout(str(tmpdir.join('newtree')),
+                                                 path_scheme=ndef_install_path_scheme)  # noqa: E501
+
+        # Install package from buildcache
+        buildcache_cmd('install', '-a', '-u', '-f', sspec.name)
+
+        # Continue blowing away caches
+        bindist.clear_spec_cache()
+        spack.stage.purge()
+
+        # test that the sbang was updated by the move
+        sbang_style_1_expected = '''{0}
+#!/usr/bin/env python
+
+{1}
+        '''.format(sbang.sbang_shebang_line(), sspec.prefix.bin)
+        sbang_style_2_expected = '''{0}
+#!/usr/bin/env python
+
+{1}
+        '''.format(sbang.sbang_shebang_line(), sspec.prefix.bin)
+
+        installed_script_style_1_path = sspec.prefix.bin.join('sbang-style-1.sh')
+        assert sbang_style_1_expected == \
+            open(str(installed_script_style_1_path)).read()
+
+        installed_script_style_2_path = sspec.prefix.bin.join('sbang-style-2.sh')
+        assert sbang_style_2_expected == \
+            open(str(installed_script_style_2_path)).read()
+
+        uninstall_cmd('-y', '/%s' % sspec.dag_hash())
+
+    finally:
+        spack.store.store = real_store
+        spack.store.layout = real_layout
diff --git a/lib/spack/spack/test/conftest.py b/lib/spack/spack/test/conftest.py
index 8615b14abe..d4f38f2c1b 100644
--- a/lib/spack/spack/test/conftest.py
+++ b/lib/spack/spack/test/conftest.py
@@ -657,6 +657,36 @@ def mock_store(tmpdir_factory, mock_repo_path, mock_configuration,
 
 
 @pytest.fixture(scope='function')
+def mutable_mock_store(tmpdir_factory, mock_repo_path, mock_configuration,
+                       _store_dir_and_cache):
+    """Creates a read-only mock database with some packages installed note
+    that the ref count for dyninst here will be 3, as it's recycled
+    across each install.
+
+    This does not actually activate the store for use by Spack -- see the
+    ``database`` fixture for that.
+
+    """
+    store_path, store_cache = _store_dir_and_cache
+    store = spack.store.Store(str(store_path))
+
+    # If the cache does not exist populate the store and create it
+    if not os.path.exists(str(store_cache.join('.spack-db'))):
+        with use_configuration(mock_configuration):
+            with use_store(store):
+                with use_repo(mock_repo_path):
+                    _populate(store.db)
+        store_path.copy(store_cache, mode=True, stat=True)
+
+    # Make the DB filesystem read-only to ensure we can't modify entries
+    store_path.join('.spack-db').chmod(mode=0o555, rec=1)
+
+    yield store
+
+    store_path.join('.spack-db').chmod(mode=0o755, rec=1)
+
+
+@pytest.fixture(scope='function')
 def database(mock_store, mock_packages, config, monkeypatch):
     """This activates the mock store, packages, AND config."""
     monkeypatch.setattr(spack.store, 'store', mock_store)
diff --git a/lib/spack/spack/test/packaging.py b/lib/spack/spack/test/packaging.py
index e2cef22577..31ad3f1a7a 100644
--- a/lib/spack/spack/test/packaging.py
+++ b/lib/spack/spack/test/packaging.py
@@ -196,10 +196,8 @@ def test_relocate_text(tmpdir):
             script.close()
         filenames = [filename]
         new_dir = '/opt/rh/devtoolset/'
-        relocate_text(filenames, old_dir, new_dir,
-                      old_dir, new_dir,
-                      old_dir, new_dir,
-                      {old_dir: new_dir})
+        # Singleton dict doesn't matter if Ordered
+        relocate_text(filenames, {old_dir: new_dir})
         with open(filename, "r")as script:
             for line in script:
                 assert(new_dir in line)
diff --git a/lib/spack/spack/test/relocate.py b/lib/spack/spack/test/relocate.py
index e883bd44c5..afeb2162db 100644
--- a/lib/spack/spack/test/relocate.py
+++ b/lib/spack/spack/test/relocate.py
@@ -12,7 +12,6 @@ import llnl.util.filesystem
 import pytest
 import spack.architecture
 import spack.concretize
-import spack.hooks.sbang as sbang
 import spack.paths
 import spack.relocate
 import spack.spec
@@ -281,7 +280,7 @@ def test_replace_prefix_bin(hello_world):
     executable = hello_world(rpaths=['/usr/lib', '/usr/lib64'])
 
     # Relocate the RPATHs
-    spack.relocate._replace_prefix_bin(str(executable), '/usr', '/foo')
+    spack.relocate._replace_prefix_bin(str(executable), {b'/usr': b'/foo'})
 
     # Some compilers add rpaths so ensure changes included in final result
     assert '/foo/lib:/foo/lib64' in rpaths_for(executable)
@@ -382,11 +381,12 @@ def test_relocate_text_bin(hello_world, copy_binary, tmpdir):
     assert not text_in_bin(str(new_binary.dirpath()), new_binary)
 
     # Check this call succeed
+    orig_path_bytes = str(orig_binary.dirpath()).encode('utf-8')
+    new_path_bytes = str(new_binary.dirpath()).encode('utf-8')
+
     spack.relocate.relocate_text_bin(
         [str(new_binary)],
-        str(orig_binary.dirpath()), str(new_binary.dirpath()),
-        spack.paths.spack_root, spack.paths.spack_root,
-        {str(orig_binary.dirpath()): str(new_binary.dirpath())}
+        {orig_path_bytes: new_path_bytes}
     )
 
     # Check original directory is not there anymore and it was
@@ -395,55 +395,13 @@ def test_relocate_text_bin(hello_world, copy_binary, tmpdir):
     assert text_in_bin(str(new_binary.dirpath()), new_binary)
 
 
-def test_relocate_text_bin_raise_if_new_prefix_is_longer():
-    short_prefix = '/short'
-    long_prefix = '/much/longer'
+def test_relocate_text_bin_raise_if_new_prefix_is_longer(tmpdir):
+    short_prefix = b'/short'
+    long_prefix = b'/much/longer'
+    fpath = str(tmpdir.join('fakebin'))
+    with open(fpath, 'w') as f:
+        f.write('/short')
     with pytest.raises(spack.relocate.BinaryTextReplaceError):
         spack.relocate.relocate_text_bin(
-            ['item'], short_prefix, long_prefix, None, None, None
+            [fpath], {short_prefix: long_prefix}
         )
-
-
-@pytest.mark.parametrize("sbang_line", [
-    "#!/bin/bash /path/to/orig/spack/bin/sbang",
-    "#!/bin/sh /orig/layout/root/bin/sbang"
-])
-def test_relocate_text_old_sbang(tmpdir, sbang_line):
-    """Ensure that old and new sbang styles are relocated."""
-
-    old_install_prefix = "/orig/layout/root/orig/install/prefix"
-    new_install_prefix = os.path.join(
-        spack.store.layout.root, "new", "install", "prefix"
-    )
-
-    # input file with an sbang line
-    original = """\
-{0}
-#!/usr/bin/env python
-
-/orig/layout/root/orig/install/prefix
-""".format(sbang_line)
-
-    # expected relocation
-    expected = """\
-{0}
-#!/usr/bin/env python
-
-{1}
-""".format(sbang.sbang_shebang_line(), new_install_prefix)
-
-    path = tmpdir.ensure("path", "to", "file")
-    with path.open("w") as f:
-        f.write(original)
-
-    spack.relocate.relocate_text(
-        [str(path)],
-        "/orig/layout/root",   spack.store.layout.root,
-        old_install_prefix,    new_install_prefix,
-        "/path/to/orig/spack", spack.paths.spack_root,
-        {
-            old_install_prefix: new_install_prefix
-        }
-    )
-
-    assert expected == open(str(path)).read()
diff --git a/lib/spack/spack/util/web.py b/lib/spack/spack/util/web.py
index 1aad85550a..1b415f5de2 100644
--- a/lib/spack/spack/util/web.py
+++ b/lib/spack/spack/util/web.py
@@ -41,6 +41,7 @@ import spack.url
 import spack.util.crypto
 import spack.util.s3 as s3_util
 import spack.util.url as url_util
+import llnl.util.lang
 
 from spack.util.compression import ALLOWED_ARCHIVE_TYPES
 
@@ -424,12 +425,6 @@ def spider(root_urls, depth=0, concurrency=32):
 
         return pages, links, subcalls
 
-    # TODO: Needed until we drop support for Python 2.X
-    def star(func):
-        def _wrapper(args):
-            return func(*args)
-        return _wrapper
-
     if isinstance(root_urls, six.string_types):
         root_urls = [root_urls]
 
@@ -450,7 +445,7 @@ def spider(root_urls, depth=0, concurrency=32):
             tty.debug("SPIDER: [depth={0}, max_depth={1}, urls={2}]".format(
                 current_depth, depth, len(spider_args))
             )
-            results = tp.map(star(_spider), spider_args)
+            results = tp.map(llnl.util.lang.star(_spider), spider_args)
             spider_args = []
             collect = current_depth < depth
             for sub_pages, sub_links, sub_spider_args in results:
diff --git a/var/spack/repos/builtin.mock/packages/old-sbang/package.py b/var/spack/repos/builtin.mock/packages/old-sbang/package.py
new file mode 100644
index 0000000000..3308f91611
--- /dev/null
+++ b/var/spack/repos/builtin.mock/packages/old-sbang/package.py
@@ -0,0 +1,35 @@
+# Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+
+from spack import *
+
+
+class OldSbang(Package):
+    """Toy package for testing the old sbang replacement problem"""
+
+    homepage = "https://www.example.com"
+    url      = "https://www.example.com/old-sbang.tar.gz"
+
+    version('1.0.0', '0123456789abcdef0123456789abcdef')
+
+    def install(self, spec, prefix):
+        mkdirp(prefix.bin)
+
+        sbang_style_1 = '''#!/bin/bash {0}/bin/sbang
+#!/usr/bin/env python
+
+{1}
+        '''.format(spack.paths.prefix, prefix.bin)
+        sbang_style_2 = '''#!/bin/sh {0}/bin/sbang
+#!/usr/bin/env python
+
+{1}
+        '''.format(spack.store.unpadded_root, prefix.bin)
+        with open('%s/sbang-style-1.sh' % self.prefix.bin, 'w') as f:
+            f.write(sbang_style_1)
+
+        with open('%s/sbang-style-2.sh' % self.prefix.bin, 'w') as f:
+            f.write(sbang_style_2)