mongodb/site_scons/site_tools/auto_install_binaries.py

# Copyright 2019 MongoDB Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# TODO: Versioned libraries
# TODO: library dependency chaining for windows dynamic builds, static dev packages
# TODO: Injectible component dependencies (jscore -> resmoke, etc.)
# TODO: Handle chmod state
# TODO: Installing resmoke and configurations
# TODO: package decomposition
# TODO: Install/package target help text
# TODO: implement sdk_headers

import os
import sys
import shlex
import itertools
from collections import defaultdict, namedtuple

import SCons
from SCons.Tool import install

ALIAS_MAP = "AIB_ALIAS_MAP"
BASE_ROLE = "AIB_BASE_ROLE"
COMPONENTS = "AIB_COMPONENTS_EXTRA"
INSTALL_ACTIONS = "AIB_INSTALL_ACTIONS"
META_ROLE = "AIB_META_ROLE"
PACKAGE_ALIAS_MAP = "AIB_PACKAGE_ALIAS_MAP"
PRIMARY_COMPONENT = "AIB_COMPONENT"
PRIMARY_ROLE = "AIB_ROLE"
ROLES = "AIB_ROLES"
ROLE_DECLARATIONS = "AIB_ROLE_DECLARATIONS"
SUFFIX_MAP = "AIB_SUFFIX_MAP"

AIB_MAKE_ARCHIVE_CONTENT = """
import os
import sys
from shutil import which

USAGE = '''
Usage: {} ARCHIVE_TYPE ARCHIVE_NAME ROOT_DIRECTORY FILES...

FILES should be absolute paths or relative to ROOT_DIRECTORY.

ARCHIVE_TYPE is one of zip or tar.
'''

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print(sys.argv[0], "takes at minimum four arguments.")
        print(USAGE.format(sys.argv[0]))
        sys.exit(1)

    archive_type = sys.argv[1]
    archive_name = sys.argv[2]
    root_dir = sys.argv[3]
    files = sys.argv[4:]

    if archive_type not in ("zip", "tar"):
        print("unsupported archive_type", archive_type)
        print(USAGE.format(sys.argv[0]))
        sys.exit(1)

    if archive_type == "tar" and which("tar") is not None:
        import subprocess
        import shlex
        tar = which("tar")
        tar_cmd = "{tar} -C {root_dir} -czf {archive_name} {files}".format(
            tar=tar,
            root_dir=root_dir,
            archive_name=archive_name,
            files=" ".join(files),
        )
        subprocess.run(shlex.split(tar_cmd))
        sys.exit(0)

    if archive_type == "zip":
        import zipfile
        archive = zipfile.ZipFile(archive_name, mode='w', compression=zipfile.ZIP_DEFLATED)
        add_file = archive.write
    else:
        import tarfile
        archive = tarfile.open(archive_name, mode='w:gz')
        add_file = archive.add

    os.chdir(root_dir)

    for filename in files:
        add_file(filename)

    archive.close()
"""

RoleInfo = namedtuple(
    "RoleInfo",
    ["alias_name", "alias", "components", "roles", "actions", "dependencies"],
)

SuffixMap = namedtuple("SuffixMap", ["directory", "default_roles",],)


class DeclaredRole:
    def __init__(self, name, dependencies=None, transitive=False, silent=False):
        self.name = name

        if dependencies is None:
            self.dependencies = set()
        else:
            self.dependencies = {dep for dep in dependencies if dep is not None}

        self.transitive = transitive
        self.silent = silent


def declare_role(env, **kwargs):
    """Construct a new role declaration"""
    return DeclaredRole(**kwargs)


def declare_roles(env, roles, base_role=None, meta_role=None):
    """Given a list of role declarations, validate them and store them in the environment"""

    role_names = [role.name for role in roles]
    if len(role_names) != len(set(role_names)):
        raise Exception("Cannot declare duplicate roles")

    # Ensure that all roles named in dependency lists actually were
    # passed in as a role.
    for role in roles:
        for d in role.dependencies:
            if d not in role_names:
                raise Exception(
                    "Role dependency '{}' does not name a declared role".format(d)
                )

    if isinstance(base_role, str):
        if base_role not in role_names:
            raise Exception(
                "A base_role argument was provided but it does not name a declared role"
            )
    elif isinstance(base_role, DeclaredRole):
        if base_role not in roles:
            raise Exception(
                "A base_role argument was provided but it is not a declared role"
            )
    elif base_role is not None:
        raise Exception(
            "The base_role argument must be a string name of a role or a role object"
        )
    else:
        # Set it to something falsy
        base_role = str()

    if isinstance(meta_role, str):
        if meta_role not in role_names:
            raise Exception(
                "A meta_role argument was provided but it does not name a declared role"
            )
    elif isinstance(meta_role, DeclaredRole):
        if meta_role not in roles:
            raise Exception(
                "A meta_role argument was provided but it is not a declared role"
            )
    elif meta_role is not None:
        raise Exception(
            "The meta_role argument must be a string name of a role or a role object"
        )
    else:
        # Set it to something falsy
        meta_role = str()

    silents = [role for role in roles if role.silent]
    if len(silents) > 1:
        raise Exception("No more than one role can be declared as silent")

    # If a base role was given, then add it as a dependency of every
    # role that isn't the base role (which would be circular).
    if base_role:
        for role in roles:
            if role.name != base_role:
                role.dependencies.add(base_role)

    # Become a dictionary, so we can look up roles easily.
    roles = {role.name: role for role in roles}

    # If a meta role was given, then add every role which isn't the
    # meta role as one of its dependencies.
    if meta_role:
        roles[meta_role].dependencies.update(r for r in roles.keys() if r != meta_role)

    # TODO: Check for DAG

    # TODO: What if base_role or meta_role is really None?
    env[BASE_ROLE] = base_role
    env[META_ROLE] = meta_role
    env[ROLE_DECLARATIONS] = roles


def generate_alias(env, component, role, target="install"):
    """Generate a scons alias for the component and role combination"""
    return "{target}-{component}{role}".format(
        target=target,
        component=component,
        role="" if env[ROLE_DECLARATIONS][role].silent else "-" + role,
    )


def get_alias_map_entry(env, component, role):
    c_entry = env[ALIAS_MAP][component]
    try:
        return c_entry[role]
    except KeyError:
        alias_name = generate_alias(env, component, role)
        r_entry = RoleInfo(
            alias_name=alias_name,
            alias=[],
            components=set(),
            roles=set(),
            actions=[],
            dependencies=[],
        )
        c_entry[role] = r_entry
        return r_entry


def get_package_name(env, component, role):
    """Return the package file name for the component and role combination."""
    basename = env[PACKAGE_ALIAS_MAP].get(
        (component, role), "{component}-{role}".format(component=component, role=role)
    )

    return basename


def get_dependent_actions(
    env, components, roles, non_transitive_roles, node, cb=None,
):
    """
    Check if node is a transitive dependency of components and roles

    If cb is not None and is callable then it will be called with all
    the arguments that get_dependent_actions was called with (except
    for cb itself) as well as the results of node_roles and the
    aib_install_actions that this function would have returned. The
    return of cb should be the dependent actions. This allows cb to
    access the results of scanning and modify the returned results via
    additional filtering.

    Returns the dependent actions.
    """
    actions = getattr(node.attributes, INSTALL_ACTIONS, None)
    if not actions:
        return []

    # Determine if the roles have any overlap with non_transitive_roles
    #
    # If they are overlapping then that means we can't transition to a
    # new role during scanning.
    if env[BASE_ROLE] not in roles:
        can_transfer = non_transitive_roles and roles.isdisjoint(non_transitive_roles)
    else:
        can_transfer = True

    node_roles = {
        role
        for role in getattr(node.attributes, ROLES, set())
        if role != env[META_ROLE]
    }
    if (
        # TODO: make the "always transitive" roles configurable
        env[BASE_ROLE] not in node_roles
        # If we are not transferrable
        and not can_transfer
        # Checks if we are actually crossing a boundry
        and node_roles.isdisjoint(roles)
    ):
        return []

    if cb is not None and callable(cb):
        return cb(components, roles, non_transitive_roles, node, node_roles, actions,)
    return actions


def scan_for_transitive_install(node, env, cb=None):
    """Walk the children of node finding all installed dependencies of it."""
    results = []
    install_sources = node.sources
    # Filter out all
    components = {
        component
        for component in getattr(node.sources[0].attributes, COMPONENTS, set())
        if component != "all"
    }
    roles = {
        role
        for role in getattr(node.sources[0].attributes, ROLES, set())
        if role != env[META_ROLE]
    }

    non_transitive_roles = {
        role for role in roles if env[ROLE_DECLARATIONS][role].transitive
    }
    for install_source in install_sources:
        install_executor = install_source.get_executor()
        if not install_executor:
            continue
        install_targets = install_executor.get_all_targets()
        if not install_targets:
            continue
        for install_target in install_targets:
            grandchildren = install_target.children()
            for grandchild in grandchildren:
                results.extend(
                    get_dependent_actions(
                        env, components, roles, non_transitive_roles, grandchild, cb=cb,
                    )
                )

    # Produce deterministic output for caching purposes
    results = sorted(results, key=str)
    return results


def collect_transitive_files(env, source, installed, cache=None):
    """Collect all installed transitive files for source where source is a list of either Alias or File nodes."""

    if not cache:
        cache = set()

    files = []

    for s in source:

        cache.add(s)

        if isinstance(s, SCons.Node.FS.File):
            if s not in installed:
                continue
            files.append(s)

        children_to_collect = []
        for child in s.children():
            if child in cache:
                continue
            if isinstance(child, SCons.Node.FS.File) and child not in installed:
                continue
            children_to_collect.append(child)

        if children_to_collect:
            files.extend(
                collect_transitive_files(env, children_to_collect, installed, cache)
            )

    return files


def archive_builder(source, target, env, for_signature):
    """Build archives of the AutoInstall'd sources."""
    if not source:
        return

    source = env.Flatten([source])
    make_archive_script = source[0].get_abspath()

    # We expect this to be a list of aliases, but really they could be
    # any sort of node.
    aliases = source[1:]

    common_ancestor = None
    archive_type = env["__AIB_ARCHIVE_TYPE"]

    # Get the path elements that make up both DESTDIR and PREFIX. Then
    # iterate the dest_dir_elems with the prefix path elements
    # stripped off the end of the path converting them to strings for
    # joining to make the common_ancestor.
    #
    # We pass the common_ancestor to tar via -C so that $PREFIX is
    # preserved in the tarball.
    dest_dir_elems = env.Dir("$DESTDIR").get_abspath()
    prefix_elems = env.subst("$PREFIX")
    # In python slicing a string with [:-0] gives an empty string. So
    # make sure we have a prefix to slice off before trying it.
    if prefix_elems:
        common_ancestor = dest_dir_elems[: -len(prefix_elems)]
    else:
        common_ancestor = dest_dir_elems

    archive_name = env.File(target[0]).get_abspath()

    command_prefix = "{python} {make_archive_script} {archive_type} {archive_name} {common_ancestor}".format(
        python=sys.executable,
        archive_type=archive_type,
        archive_name=archive_name,
        make_archive_script=make_archive_script,
        common_ancestor=common_ancestor,
    )

    # If we are just being invoked for our signature, we can omit the indirect dependencies
    # found by expanding the transitive dependencies, since we really only have a hard dependency
    # on our direct depenedencies.
    if for_signature:
        return command_prefix

    # Pre-process what should be in the archive. We need to pass the
    # set of known installed files along to the transitive dependency
    # walk so we can filter out files that aren't in the install
    # directory.
    installed = env.get("__AIB_INSTALLED_SET", set())
    transitive_files = collect_transitive_files(env, aliases, installed)
    paths = {file.get_abspath() for file in transitive_files}

    # The env["ESCAPE"] function is used by scons to make arguments
    # valid for the platform that we're running on. For instance it
    # will properly quote paths that have spaces in them on Posix
    # platforms and handle \ / on Windows.
    escape_func = env.get("ESCAPE", lambda x: x)

    # TODO: relpath is costly, and we do it for every file in the archive here. We should
    # find a way to avoid the repeated relpath invocation, probably by bucketing by directory.
    relative_files = " ".join(
        [escape_func(os.path.relpath(path, common_ancestor)) for path in paths]
    )

    return " ".join([command_prefix, relative_files])


def auto_install(env, target, source, **kwargs):
    """Auto install builder."""
    source = [env.Entry(s) for s in env.Flatten([source])]
    roles = {
        kwargs.get(PRIMARY_ROLE),
    }

    if env[META_ROLE]:
        roles.add(env[META_ROLE])

    if kwargs.get(ROLES) is not None:
        roles = roles.union(set(kwargs[ROLES]))

    component = kwargs.get(PRIMARY_COMPONENT)
    if component is not None and (not isinstance(component, str) or " " in component):
        raise Exception("AIB_COMPONENT must be a string and contain no whitespace.")

    components = {
        component,
        # The 'all' tag is implicitly attached as a component
        "all",
    }
    # Some tools will need to create multiple components so we add
    # this "hidden" argument that accepts a set or list.
    #
    # Use get here to check for existence because it is rarely
    # ommitted as a kwarg (because it is set by the default emitter
    # for all common builders), but is often set to None.
    if kwargs.get(COMPONENTS) is not None:
        components = components.union(set(kwargs[COMPONENTS]))

    # Remove false values such as None or ""
    roles = {role for role in roles if role}
    components = {component for component in components if component}

    actions = []

    for s in source:
        s.attributes.keep_targetinfo = 1
        setattr(s.attributes, COMPONENTS, components)
        setattr(s.attributes, ROLES, roles)

        # We must do an eearly subst here so that the _aib_debugdir
        # generator has a chance to run while seeing 'source'.
        #
        # TODO: Find a way to not need this early subst.
        target = env.Dir(env.subst(target, source=source))

        action = env.Install(target=target, source=s,)

        setattr(
            s.attributes,
            INSTALL_ACTIONS,
            action if isinstance(action, (list, set)) else [action],
        )
        actions.append(action)

    actions = env.Flatten(actions)
    for component, role in itertools.product(components, roles):

        entry = get_alias_map_entry(env, component, role)
        entry.components.update(components)
        entry.roles.update(roles)
        entry.actions.extend(actions)

        # TODO: this hard codes behavior that should be done configurably
        if component != "common":
            dentry = get_alias_map_entry(env, "common", role)
            entry.dependencies.append(dentry)

    return actions


def finalize_install_dependencies(env):
    """Generates package aliases and wires install dependencies."""

    installed = set(env.FindInstalledFiles())

    for component, rolemap in env[ALIAS_MAP].items():
        for role, info in rolemap.items():
            info.alias.extend(env.Alias(info.alias_name, info.actions))
            setattr(info.alias[0].attributes, COMPONENTS, info.components)
            setattr(info.alias[0].attributes, ROLES, info.roles)
            env.Depends(info.alias, [d.alias for d in info.dependencies])

    common_rolemap = env[ALIAS_MAP].get("common")
    default_rolemap = env[ALIAS_MAP].get("default")

    if default_rolemap and "runtime" in default_rolemap:
        env.Alias("install", "install-default")
        env.Default("install")

    # TODO: $BUILD_ROOT should be $VARIANT_DIR after we fix our dir
    # setup later on.
    make_archive_script = env.Textfile(
        target="$BUILD_ROOT/aib_make_archive.py", source=[AIB_MAKE_ARCHIVE_CONTENT],
    )

    for component, rolemap in env[ALIAS_MAP].items():
        for role, info in rolemap.items():

            if common_rolemap and component != "common" and role in common_rolemap:
                env.Depends(info.alias, common_rolemap[role].alias)

            role_decl = env[ROLE_DECLARATIONS].get(role)
            for dependency in role_decl.dependencies:
                dependency_info = rolemap.get(dependency, [])
                if dependency_info:
                    env.Depends(info.alias, dependency_info.alias)

            pkg_name = get_package_name(env, component, role)

            for fmt in ("zip", "tar"):
                # TODO: $PKGDIR support
                if fmt == "zip":
                    pkg_suffix = "$AIB_ZIP_SUFFIX"
                else:
                    pkg_suffix = "$AIB_TARBALL_SUFFIX"

                archive = env.__AibArchive(
                    target="#{}.{}".format(pkg_name, pkg_suffix),
                    source=[make_archive_script] + info.alias,
                    __AIB_ARCHIVE_TYPE=fmt,
                    __AIB_INSTALLED_SET=installed,
                    AIB_COMPONENT=component,
                    AIB_ROLE=role,
                )

                # TODO: perhaps caching of packages / tarballs should be
                # configurable? It's possible someone would want to do it.
                env.NoCache(archive)

                archive_alias = generate_alias(env, component, role, target=fmt)
                env.Alias(archive_alias, archive)


def auto_install_emitter(target, source, env):
    """When attached to a builder adds an appropriate AutoInstall to that Builder."""
    for t in target:
        entry = env.Entry(t)
        suffix = entry.get_suffix()
        if env.get("AIB_IGNORE", False):
            continue

        # There is no API for determining if an Entry is operating in
        # a SConf context. We obviously do not want to auto tag, and
        # install conftest Programs. So we filter them out the only
        # way available to us.
        #
        # We're working with upstream to expose this information.
        if "conftest" in str(entry):
            continue

        auto_install_mapping = env[SUFFIX_MAP].get(suffix)
        if auto_install_mapping is not None:
            env.AutoInstall(
                auto_install_mapping.directory,
                entry,
                AIB_COMPONENT=env.get(PRIMARY_COMPONENT),
                AIB_ROLE=env.get(PRIMARY_ROLE),
                AIB_ROLES=auto_install_mapping.default_roles,
                AIB_COMPONENTS_EXTRA=env.get(COMPONENTS),
            )

    return (target, source)


def add_suffix_mapping(env, suffix, role=None):
    """Map suffix to role"""
    if isinstance(suffix, str):
        if role not in env[ROLE_DECLARATIONS]:
            raise Exception(
                "target {} is not a known role. Available roles are {}".format(
                    role, env[ROLE_DECLARATIONS].keys()
                )
            )
        env[SUFFIX_MAP][env.subst(suffix)] = role

    if not isinstance(suffix, dict):
        raise Exception("source must be a dictionary or a string")

    for _, mapping in suffix.items():
        for role in mapping.default_roles:
            if role not in env[ROLE_DECLARATIONS]:
                raise Exception(
                    "target {} is not a known role. Available roles are {}".format(
                        target, env[ROLE_DECLARATIONS].keys()
                    )
                )

    env[SUFFIX_MAP].update({env.subst(key): value for key, value in suffix.items()})


def add_package_name_alias(env, component, role, name):
    """Add a package name mapping for the combination of component and role."""
    # Verify we didn't get a None or empty string for any argument
    if not name:
        raise Exception(
            "when setting a package name alias must provide a name parameter"
        )
    if not component:
        raise Exception("No component provided for package name alias")
    if not role:
        raise Exception("No role provided for package name alias")
    env[PACKAGE_ALIAS_MAP][(component, role)] = name


def suffix_mapping(env, directory=False, default_roles=False):
    """Generate a SuffixMap object from source and target."""
    return SuffixMap(directory=directory, default_roles=default_roles,)


def dest_dir_generator(initial_value=None):
    """Memoized dest_dir_generator"""
    dd = (None, None)

    def generator(source, target, env, for_signature):
        nonlocal dd

        # SCons does not perform substitution for "sub" Dir calls on a
        # Dir Node. Additionally we need to determine if it's an
        # absolute path here because if it is the sub Dir call will
        # not expand correctly.
        prefix = env.subst("$PREFIX")
        if prefix and prefix[0] == "/":
            prefix = prefix[1:]

        if dd[1] is not None and dd[0] == prefix:
            return dd[1]

        if initial_value is None:
            dest_dir = env.Dir("#install")
        elif isinstance(initial_value, str):
            dest_dir = env.Dir(initial_value)
        elif isinstance(initial_value, SCons.Node.FS.Dir):
            dest_dir = initial_value
        else:
            raise Exception("initial DESTDIR value must be string or Dir")

        dd = (prefix, dest_dir.Dir(prefix))
        return dd[1]

    return generator


def _aib_debugdir(source, target, env, for_signature):
    for s in source:
        # TODO: We shouldn't need to reach into the attributes of the debug tool like this.
        origin = getattr(s.attributes, "debug_file_for", None)
        oentry = env.Entry(origin)
        osuf = oentry.get_suffix()
        return env[SUFFIX_MAP].get(osuf)[0]


def exists(_env):
    """Always activate this tool."""
    return True


def list_components(env, **kwargs):
    """List registered components for env."""
    print("Known AIB components:")
    for key in env[ALIAS_MAP]:
        print("\t", key)


def list_targets(env, **kwargs):
    """List AIB generated targets for env."""
    print("Generated AIB targets:")
    for _, rolemap in env[ALIAS_MAP].items():
        for _, info in rolemap.items():
            print("\t", info.alias[0].name)


def generate(env):  # pylint: disable=too-many-statements
    """Generate the auto install builders."""
    bld = SCons.Builder.Builder(
        action=SCons.Action.CommandGeneratorAction(
            archive_builder,
            {"cmdstr": "Building package ${TARGETS[0]} from ${SOURCES[1:]}"},
        )
    )
    env.Append(BUILDERS={"__AibArchive": bld})
    env["AIB_TARBALL_SUFFIX"] = env.get("AIB_TARBALL_SUFFIX", "tar.gz")
    env["AIB_ZIP_SUFFIX"] = env.get("AIB_ZIP_SUFFIX", "zip")

    # Matches the autoconf documentation:
    # https://www.gnu.org/prep/standards/html_node/Directory-Variables.html
    env["DESTDIR"] = dest_dir_generator(env.get("DESTDIR", None))
    env["PREFIX_BINDIR"] = env.get("PREFIX_BINDIR", "$DESTDIR/bin")
    env["PREFIX_LIBDIR"] = env.get("PREFIX_LIBDIR", "$DESTDIR/lib")
    env["PREFIX_SHAREDIR"] = env.get("PREFIX_SHAREDIR", "$DESTDIR/share")
    env["PREFIX_DOCDIR"] = env.get("PREFIX_DOCDIR", "$PREFIX_SHAREDIR/doc")
    env["PREFIX_INCLUDEDIR"] = env.get("PREFIX_INCLUDEDIR", "$DESTDIR/include")
    env["PREFIX_DEBUGDIR"] = env.get("PREFIX_DEBUGDIR", _aib_debugdir)
    env[SUFFIX_MAP] = {}
    env[PACKAGE_ALIAS_MAP] = {}
    env[ALIAS_MAP] = defaultdict(dict)

    env.AddMethod(suffix_mapping, "SuffixMap")
    env.AddMethod(add_suffix_mapping, "AddSuffixMapping")
    env.AddMethod(add_package_name_alias, "AddPackageNameAlias")
    env.AddMethod(auto_install, "AutoInstall")
    env.AddMethod(finalize_install_dependencies, "FinalizeInstallDependencies")
    env.AddMethod(declare_role, "Role")
    env.AddMethod(declare_roles, "DeclareRoles")
    env.Tool("install")

    # TODO: we should probably expose these as PseudoBuilders and let
    # users define their own aliases for them.
    env.Alias("list-aib-components", [], [list_components])
    env.AlwaysBuild("list-aib-components")

    env.Alias("list-aib-targets", [], [list_targets])
    env.AlwaysBuild("list-aib-targets")

    for builder in ["Program", "SharedLibrary", "LoadableModule", "StaticLibrary"]:
        builder = env["BUILDERS"][builder]
        base_emitter = builder.emitter
        # TODO: investigate if using a ListEmitter here can cause
        # problems if AIB is not loaded last
        new_emitter = SCons.Builder.ListEmitter([base_emitter, auto_install_emitter,])
        builder.emitter = new_emitter

    base_install_builder = install.BaseInstallBuilder
    assert base_install_builder.target_scanner is None

    base_install_builder.target_scanner = SCons.Scanner.Scanner(
        function=scan_for_transitive_install, path_function=None
    )