Write TSV, JSON and text files atomically (#62)

save_json(), write_python(), write_text(), write_tsv() and _write_tsv_simple()
all opened the destination in 'w' mode, which truncates it before a single byte
of the new content is written. A crash, a full disk or a kill in that window
leaves the file empty or half written and the previous contents are gone.

cluster_group.tsv is the one that hurts, since it holds manual curation labels
that cannot be recomputed.

Add _atomic_open(), a context manager that writes to a temporary file in the
destination directory and moves it into place with os.replace() on a clean
exit, and route all five writers through it. The temporary file is given the
permissions of the file it replaces, or the permissions open() would have used
for a new file, since tempfile creates its files with 0600.

Co-authored-by: Cyrille Rossant <rossant@users.noreply.github.com>

Author: adityasingh2400

CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Development checkouts now report a development version with a Git commit suffix.
 - Template datasets without waveform templates can be reloaded after cluster assignments change,
   and stored spike-waveform subsets provide their waveform sample count.
+- TSV, JSON, text and `params.py` files are written atomically, so a crash during a save no longer
+  truncates the file it was replacing, for instance `cluster_group.tsv`.
 - #57 a blank `dat_path` in `params.py` is now read as "no raw data file" instead of resolving to
   the dataset directory.
 

phylib/utils/_misc.py

@@ -8,12 +8,15 @@
 #------------------------------------------------------------------------------
 
 import base64
+from contextlib import contextmanager
 import csv
 from importlib import import_module
 import json
 import logging
 import os
 from pathlib import Path
+import secrets
+import stat
 import subprocess
 from textwrap import dedent
 
@@ -24,6 +27,62 @@
 logger = logging.getLogger(__name__)
 
 
+#------------------------------------------------------------------------------
+# Atomic file writing
+#------------------------------------------------------------------------------
+
+def _copy_file_mode(src, dst):
+    """Give `src` the permissions of `dst` when `dst` exists."""
+    try:
+        mode = stat.S_IMODE(os.stat(dst).st_mode)
+    except FileNotFoundError:
+        # The exclusive open used for a new temporary file already applied the process umask.
+        return
+    os.chmod(src, mode)
+
+
+@contextmanager
+def _atomic_open(path, mode='w', **kwargs):
+    """Open a file for writing without ever truncating the file it replaces.
+
+    The contents are written to a temporary file in the same directory as the destination, and
+    that file is moved into place with `os.replace()` once the `with` block completes, which is
+    atomic as long as both paths are on the same filesystem. A reader therefore sees either the
+    old file or the new one, never a half-written one. If the block raises, the temporary file
+    is removed and any existing file is left untouched.
+
+    """
+    path = Path(path)
+    ensure_dir_exists(path.parent)
+    if not mode.startswith('w'):
+        raise ValueError("_atomic_open() only supports write modes")
+    exclusive_mode = 'x' + mode[1:]
+    # NOTE: the temporary file is hidden and does not end in a data file extension, so that a
+    # file left behind by a hard kill is not picked up by the globs used to discover datasets.
+    # Opening it with `x` both prevents a name collision and lets the OS apply the current umask,
+    # without temporarily changing that process-global setting.
+    while True:
+        tmp = path.parent / ('.%s.%s.tmp' % (path.name, secrets.token_hex(16)))
+        try:
+            f = open(tmp, exclusive_mode, **kwargs)
+            break
+        except FileExistsError:  # pragma: no cover - cryptographically unlikely
+            continue
+    try:
+        with f:
+            yield f
+            f.flush()
+            os.fsync(f.fileno())
+        _copy_file_mode(tmp, path)
+        os.replace(tmp, path)
+    except BaseException:
+        try:
+            os.remove(tmp)
+        except OSError:  # pragma: no cover
+            pass
+        raise
+
+
 #------------------------------------------------------------------------------
 # JSON utility functions
 #------------------------------------------------------------------------------
@@ -128,9 +187,7 @@ def save_json(path, data):
     """
     assert isinstance(data, dict)
     data = _stringify_keys(data)
-    path = Path(path)
-    ensure_dir_exists(path.parent)
-    with path.open('w') as f:
+    with _atomic_open(path) as f:
         json.dump(data, f, cls=_CustomEncoder, indent=2, sort_keys=True)
 
 
@@ -190,7 +247,7 @@ def write_python(path, data):
     -------
 
     """
-    with open(path, 'w') as f:
+    with _atomic_open(path) as f:
         for k, v in data.items():
             if isinstance(v, str):
                 v = '"%s"' % v
@@ -206,9 +263,8 @@ def read_text(path):
 def write_text(path, contents):
     """Write a text file."""
     contents = dedent(contents)
-    path = Path(path)
-    ensure_dir_exists(path.parent)
-    path.write_text(contents)
+    with _atomic_open(path) as f:
+        f.write(contents)
 
 
 def _try_make_number(value):
@@ -266,9 +322,8 @@ def write_tsv(path, data, first_field=None, exclude_fields=(), n_significant_fig
 
     """
     path = Path(path)
-    ensure_dir_exists(path.parent)
     delimiter = '\t' if path.suffix == '.tsv' else ','
-    with path.open('w', newline='') as f:
+    with _atomic_open(path, newline='') as f:
         if not data:
             logger.info("Data was empty when writing %s.", path)
             return
@@ -327,9 +382,8 @@ def _write_tsv_simple(path, field_name, data):
 
     """
     path = Path(path)
-    ensure_dir_exists(path.parent)
     delimiter = '\t' if path.suffix == '.tsv' else ','
-    with path.open('w', newline='') as f:
+    with _atomic_open(path, newline='') as f:
         writer = csv.writer(f, delimiter=delimiter)
         writer.writerow(['cluster_id', field_name])
         writer.writerows([(cluster_id, data[cluster_id]) for cluster_id in sorted(data)])

phylib/utils/tests/test_misc.py

@@ -6,15 +6,19 @@
 # Imports
 #------------------------------------------------------------------------------
 
+import csv
+import json
+import os
 import re
+import stat
 
 import numpy as np
 from numpy.testing import assert_array_equal as ae
 from pytest import raises, mark
 
 from .._misc import (
-    _git_version, load_json, save_json, load_pickle, save_pickle, read_python, write_python,
-    read_text, write_text, _read_tsv_simple, _write_tsv_simple, read_tsv, write_tsv,
+    _atomic_open, _git_version, load_json, save_json, load_pickle, save_pickle, read_python,
+    write_python, read_text, write_text, _read_tsv_simple, _write_tsv_simple, read_tsv, write_tsv,
     _pretty_floats, _encode_qbytearray, _decode_qbytearray, _fullname, _load_from_fullname)
 
 
@@ -137,6 +141,100 @@ def test_write_tsv(tempdir):
     assert read_tsv(path) == data
 
 
+class _Boom(Exception):
+    pass
+
+
+def test_atomic_open(tempdir):
+    path = tempdir / 'test_dir/test.txt'
+
+    with _atomic_open(path) as f:
+        f.write('hello')
+    assert read_text(path) == 'hello'
+    # The temporary file must not be left behind.
+    assert sorted(p.name for p in path.parent.iterdir()) == ['test.txt']
+
+    # A new file gets the permissions open() would have given it, without reading or changing
+    # the process-global umask. Compare it with an ordinary file created in the same directory.
+    control = path.parent / 'control.txt'
+    control.write_text('control')
+    assert stat.S_IMODE(path.stat().st_mode) == stat.S_IMODE(control.stat().st_mode)
+    control.unlink()
+
+    # The permissions of an existing file are preserved.
+    expected_mode = stat.S_IMODE(path.stat().st_mode)
+    if os.name != 'nt':
+        path.chmod(0o640)
+        expected_mode = 0o640
+    with _atomic_open(path) as f:
+        f.write('world')
+    assert read_text(path) == 'world'
+    assert stat.S_IMODE(path.stat().st_mode) == expected_mode
+
+    # A failure half-way through leaves the previous file untouched and no leftovers.
+    with raises(_Boom):
+        with _atomic_open(path) as f:
+            f.write('this should never be visible')
+            raise _Boom()
+    assert read_text(path) == 'world'
+    assert sorted(p.name for p in path.parent.iterdir()) == ['test.txt']
+
+    # Binary writes use the same atomic and cleanup path, ready for numpy saves.
+    with _atomic_open(path, mode='wb') as f:
+        f.write(b'bytes')
+    assert path.read_bytes() == b'bytes'
+
+
+@mark.parametrize('writer', ['write_tsv', '_write_tsv_simple'])
+def test_write_tsv_atomic(tempdir, monkeypatch, writer):
+    # Curation data such as cluster_group.tsv must survive a crash during a save.
+    path = tempdir / 'cluster_group.tsv'
+    if writer == 'write_tsv':
+        def _write(value):
+            write_tsv(path, [{'cluster_id': 2, 'group': value}], first_field='cluster_id')
+    else:
+        def _write(value):
+            _write_tsv_simple(path, 'group', {2: value})
+
+    _write('good')
+    before = path.read_bytes()
+    assert before
+    files_before = sorted(p.name for p in tempdir.iterdir())
+
+    real_writer = csv.writer
+
+    def _failing_writer(f, **kwargs):
+        # Mimic a crash after part of the file has been written.
+        real_writer(f, **kwargs).writerow(['cluster_id', 'group'])
+        raise _Boom()
+
+    monkeypatch.setattr(csv, 'writer', _failing_writer)
+    with raises(_Boom):
+        _write('mua')
+
+    assert path.read_bytes() == before
+    assert sorted(p.name for p in tempdir.iterdir()) == files_before
+
+
+def test_save_json_atomic(tempdir, monkeypatch):
+    path = tempdir / 'test.json'
+    save_json(path, {'a': 1})
+    before = path.read_bytes()
+    assert before
+    files_before = sorted(p.name for p in tempdir.iterdir())
+
+    def _failing_dump(data, f, **kwargs):
+        f.write('{"a":')
+        raise _Boom()
+
+    monkeypatch.setattr(json, 'dump', _failing_dump)
+    with raises(_Boom):
+        save_json(path, {'a': 2})
+
+    assert path.read_bytes() == before
+    assert sorted(p.name for p in tempdir.iterdir()) == files_before
+
+
 def test_git_version():
     v = _git_version()
     assert re.match(r'^\+git\.[0-9a-f]{8}(?:\.dirty)?$', v)