def populate_temp_dir_and_rename()

in antlir/fs_utils.py [0:0]


def populate_temp_dir_and_rename(dest_path, *, overwrite: bool = False) -> Path:
    """
    Returns a Path to a temporary directory. The context block may populate
    this directory, which will then be renamed to `dest_path`, optionally
    deleting any preexisting directory (if `overwrite=True`).

    If the context block throws, the partially populated temporary directory
    is removed, while `dest_path` is left alone.

    By writing to a brand-new temporary directory before renaming, we avoid
    the problems of partially writing files, or overwriting some files but
    not others.  Moreover, populate-temporary-and-rename is robust to
    concurrent writers, and tends to work on broken NFSes unlike `flock`.
    """
    dest_path = os.path.normpath(dest_path)  # Trailing / breaks `os.rename()`
    # Putting the temporary directory as a sibling minimizes permissions
    # issues, and maximizes the chance that we're on the same filesystem
    base_dir = os.path.dirname(dest_path)
    td = tempfile.mkdtemp(dir=base_dir)
    try:
        # pyre-fixme[7]: Expected `Path` but got `Generator[Path, None, None]`.
        yield Path(td)

        # Delete+rename is racy, but EdenFS lacks RENAME_EXCHANGE (t34057927)
        # Retry if we raced with another writer -- i.e., last-to-run wins.
        while True:
            if overwrite and os.path.isdir(dest_path):
                with tempfile.TemporaryDirectory(dir=base_dir) as del_dir:
                    try:
                        os.rename(dest_path, del_dir)
                    except FileNotFoundError:  # pragma: no cover
                        continue  # retry, another writer deleted first?
            try:
                os.rename(td, dest_path)
            except OSError as ex:
                if not (
                    overwrite
                    and ex.errno
                    in [
                        # Different kernels have different error codes when the
                        # target already exists and is a nonempty directory.
                        errno.ENOTEMPTY,
                        errno.EEXIST,
                    ]
                ):
                    raise
                log.exception(  # pragma: no cover
                    f"Retrying deleting {dest_path}, another writer raced us"
                )
            # We won the race
            break  # pragma: no cover
    except BaseException:
        shutil.rmtree(td)
        raise