utils.py: Add save_file_atomic() helper

This is a context manager that can be used to divert file writes into a
temporary file, which is then renamed into place once writing is
complete. It is primarily intended for use by source plugins which
download files, so they can ensure that their downloads appear atomic
and there is no risk of leaving half-downloaded files in the cache.

So far this is not used in the core, but it is needed by the Docker
source plugin that is proposed for the bst-external plugins repo. See:
https://gitlab.com/BuildStream/bst-external/merge_requests/9

3 files changed, 120 insertions, 0 deletions

diff --git a/buildstream/utils.py b/buildstream/utils.py
index 89c4cc016..9c65d8e74 100644
--- a/buildstream/utils.py
+++ b/buildstream/utils.py
@@ -452,6 +452,64 @@ def get_bst_version():
     return (int(versions[0]), int(versions[1]))
 
 
+@contextmanager
+def save_file_atomic(filename, mode='w', *, buffering=-1, encoding=None,
+                     errors=None, newline=None, closefd=True, opener=None):
+    """Save a file with a temporary name and rename it into place when ready.
+
+    This is a context manager which is meant for saving data to files.
+    The data is written to a temporary file, which gets renamed to the target
+    name when the context is closed. This avoids readers of the file from
+    getting an incomplete file.
+
+    **Example:**
+
+    .. code:: python
+
+      with save_file_atomic('/path/to/foo', 'w') as f:
+          f.write(stuff)
+
+    The file will be called something like ``tmpCAFEBEEF`` until the
+    context block ends, at which point it gets renamed to ``foo``. The
+    temporary file will be created in the same directory as the output file.
+    The ``filename`` parameter must be an absolute path.
+
+    If an exception occurs or the process is terminated, the temporary file will
+    be deleted.
+    """
+    # This feature has been proposed for upstream Python in the past, e.g.:
+    # https://bugs.python.org/issue8604
+
+    assert os.path.isabs(filename), "The utils.save_file_atomic() parameter ``filename`` must be an absolute path"
+    dirname = os.path.dirname(filename)
+    fd, tempname = tempfile.mkstemp(dir=dirname)
+    os.close(fd)
+
+    f = open(tempname, mode=mode, buffering=buffering, encoding=encoding,
+             errors=errors, newline=newline, closefd=closefd, opener=opener)
+
+    def cleanup_tempfile():
+        f.close()
+        try:
+            os.remove(tempname)
+        except FileNotFoundError:
+            pass
+        except OSError as e:
+            raise UtilError("Failed to cleanup temporary file {}: {}".format(tempname, e)) from e
+
+    try:
+        with _signals.terminator(cleanup_tempfile):
+            f.real_filename = filename
+            yield f
+            f.close()
+            # This operation is atomic, at least on platforms we care about:
+            # https://bugs.python.org/issue8828
+            os.replace(tempname, filename)
+    except Exception as e:
+        cleanup_tempfile()
+        raise
+
+
 # Recursively remove directories, ignoring file permissions as much as
 # possible.
 def _force_rmtree(rootpath, **kwargs):
diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
--- /dev/null
+++ b/tests/utils/__init__.py
diff --git a/tests/utils/savefile.py b/tests/utils/savefile.py
new file mode 100644
index 000000000..87f9f4b0a
--- /dev/null
+++ b/tests/utils/savefile.py
@@ -0,0 +1,62 @@
+import os
+import pytest
+
+from buildstream.utils import save_file_atomic
+
+
+def test_save_new_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-success.test')
+    with save_file_atomic(filename, 'w') as f:
+        f.write('foo\n')
+
+    assert os.listdir(tmpdir) == ['savefile-success.test']
+    with open(filename) as f:
+        assert f.read() == 'foo\n'
+
+
+def test_save_over_existing_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-overwrite.test')
+
+    with open(filename, 'w') as f:
+        f.write('existing contents\n')
+
+    with save_file_atomic(filename, 'w') as f:
+        f.write('overwritten contents\n')
+
+    assert os.listdir(tmpdir) == ['savefile-overwrite.test']
+    with open(filename) as f:
+        assert f.read() == 'overwritten contents\n'
+
+
+def test_exception_new_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-exception.test')
+
+    with pytest.raises(RuntimeError):
+        with save_file_atomic(filename, 'w') as f:
+            f.write('Some junk\n')
+            raise RuntimeError("Something goes wrong")
+
+    assert os.listdir(tmpdir) == []
+
+
+def test_exception_existing_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-existing.test')
+
+    with open(filename, 'w') as f:
+        f.write('existing contents\n')
+
+    with pytest.raises(RuntimeError):
+        with save_file_atomic(filename, 'w') as f:
+            f.write('Some junk\n')
+            raise RuntimeError("Something goes wrong")
+
+    assert os.listdir(tmpdir) == ['savefile-existing.test']
+    with open(filename) as f:
+        assert f.read() == 'existing contents\n'
+
+
+def test_attributes(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-attributes.test')
+    with save_file_atomic(filename, 'w') as f:
+        assert f.real_filename == filename
+        assert f.name != filename