Make save raise on failure; let save move file-backed data into the output instead of copying it

exir/_serialize/_cord.py

@@ -8,13 +8,31 @@
 import io
 import os
 import shutil
+import stat
+import sys
 import tempfile
 import weakref
-from typing import List, Optional, Union
+from typing import Container, Dict, IO, List, Optional, Union
 
 
 class FileBackedData:
-    """A byte buffer that stays on disk until explicitly closed."""
+    """A byte buffer that stays on disk until explicitly closed or moved.
+
+    ``ExecutorchProgramManager.save(path, consume=True)`` moves data that owns its file
+    (``cleanup=True``, as from ``move_from``) into the file it writes instead of copying
+    it. Chunks are copied from the end back to the start, and the file is truncated after
+    each chunk is written and synced to the output; the part of the output not yet
+    written stays a hole. On a file system with sparse files, the data is then on disk
+    once, plus one chunk; on one without, the move needs as much space as a copy. The
+    syncs make a move several times slower than a copy. The move needs a regular file
+    that is not a link and has no other name, on the output's file system. Anything else,
+    and Windows (where files are not sparse by default), gets a copy. After a move the
+    file is empty, and reading this or any other ``FileBackedData`` on it raises.
+
+    A move that fails before its first truncation leaves the data intact. One that fails
+    later leaves the file in place and says, in a note on the error, which bytes are
+    still in it and where the rest is in the output.
+    """
 
     _COPY_CHUNK_SIZE = 8 * 1024 * 1024
 
@@ -56,27 +74,166 @@ def move_from(cls, path: str, sha256: Optional[bytes] = None) -> "FileBackedData
     def __len__(self) -> int:
         return self._size
 
+    def __deepcopy__(self, memo: Dict[int, object]) -> "FileBackedData":
+        # One object per file: of two that owned it, the first collected would delete the
+        # file under the other, and one that doesn't own it can't move it.
+        return self
+
+    @property
+    def path(self) -> str:
+        return self._path
+
+    def _check_size(self, size: int) -> None:
+        # The size, not a flag, shows a moved file, so every object on it sees it.
+        if size != self._size:
+            raise RuntimeError(
+                f"{self._path} has {size} bytes, expected {self._size}; it changed after "
+                "it was wrapped, for example because a consuming save moved its data"
+            )
+
+    def check(self) -> None:
+        """Raise if the file no longer holds the data, as after a move."""
+        self._check_size(os.stat(self._path).st_size)
+
+    def _open(self) -> IO[bytes]:
+        f = open(self._path, "rb")
+        try:
+            self._check_size(os.fstat(f.fileno()).st_size)
+        except BaseException:
+            f.close()
+            raise
+        return f
+
     def prefix(self, size: int) -> bytes:
-        with open(self._path, "rb") as f:
+        with self._open() as f:
             return f.read(size)
 
     def sha256(self) -> bytes:
         if self._sha256 is None:
             digest = hashlib.sha256()
-            with open(self._path, "rb") as f:
+            with self._open() as f:
                 while chunk := f.read(self._COPY_CHUNK_SIZE):
                     digest.update(chunk)
             self._sha256 = digest.digest()
         return self._sha256
 
     def to_bytes(self) -> bytes:
-        with open(self._path, "rb") as f:
+        with self._open() as f:
             return f.read()
 
     def write_to_file(self, outfile: io.BufferedIOBase) -> None:
-        with open(self._path, "rb") as f:
+        with self._open() as f:
             shutil.copyfileobj(f, outfile, length=self._COPY_CHUNK_SIZE)
 
+    def _move_to(self, outfile: io.BufferedIOBase) -> Optional[int]:
+        """Move the data into ``outfile``, a regular file that its caller created for
+        this write (see the class docstring). Returns the offset in ``outfile`` it was
+        moved to, or None, having written nothing, if it has to be copied instead."""
+        source = self._open_to_move(outfile)
+        if source is None:
+            return None
+        try:
+            return self._move_chunks(outfile, source)
+        finally:
+            os.close(source)
+
+    def _copy_reason(
+        self, output_device: int, source: Optional[os.stat_result] = None
+    ) -> Optional[str]:
+        """Why a move into a file on ``output_device`` would copy the data instead, or
+        None if it can move. Judges ``source``, the stat of the opened file, if given, and
+        else the file at the path."""
+        if self._finalizer is None or not self._finalizer.alive:
+            return "not owned"
+        if sys.platform == "win32":
+            return "on Windows"
+        if source is None:
+            try:
+                source = os.lstat(self._path)
+            except OSError:
+                return "unreadable"
+            if not stat.S_ISREG(source.st_mode) or not os.access(self._path, os.W_OK):
+                return "a link or read-only"
+        # Truncating changes every name of the file, and on another file system it frees
+        # nothing for the output.
+        if source.st_nlink != 1:
+            return "with another hard link"
+        if source.st_dev != output_device:
+            return "on another file system"
+        if source.st_size != self._size:
+            return "changed in size"
+        return None
+
+    def _open_to_move(self, outfile: io.BufferedIOBase) -> Optional[int]:
+        output_stat = os.fstat(outfile.fileno())
+        if self._copy_reason(output_stat.st_dev) is not None:
+            return None
+        try:
+            # O_NOFOLLOW refuses a link, and the checks below look at the file opened.
+            source = os.open(self._path, os.O_RDWR | os.O_NOFOLLOW)
+        except OSError:
+            return None
+        try:
+            source_stat = os.fstat(source)
+            if os.path.samestat(source_stat, output_stat):
+                raise ValueError(f"Can't move {self._path} into itself")
+        except BaseException:
+            os.close(source)
+            raise
+        if self._copy_reason(output_stat.st_dev, source_stat) is not None:
+            os.close(source)
+            return None
+        return source
+
+    def _move_chunks(self, outfile: io.BufferedIOBase, source: int) -> int:
+        start = outfile.tell()
+        end = self._size
+        detached = False
+        try:
+            while end > 0:
+                begin = max(end - self._COPY_CHUNK_SIZE, 0)
+                chunk = os.pread(source, end - begin, begin)
+                if len(chunk) != end - begin:
+                    raise OSError(
+                        f"read {len(chunk)} of {end - begin} bytes at offset {begin} "
+                        f"of {self._path}"
+                    )
+                outfile.seek(start + begin)
+                written = outfile.write(chunk)
+                if written != len(chunk):
+                    raise OSError(
+                        f"wrote {written} of {len(chunk)} bytes at offset "
+                        f"{start + begin}"
+                    )
+                # The chunk has to be on the output's disk before the source gives it up.
+                outfile.flush()
+                os.fsync(outfile.fileno())
+                if not detached and self._finalizer is not None:
+                    # From the first truncation, part of the data is found only through
+                    # the error note, so no interrupt may let the finalizer delete it.
+                    detached = self._finalizer.detach() is not None
+                end = begin
+                os.ftruncate(source, begin)
+            if detached:
+                self._finalizer = weakref.finalize(self, self._remove, self._path)
+            outfile.seek(start + self._size)
+        except BaseException as e:
+            try:
+                # The truncation itself may have failed and left more bytes in the file.
+                end = os.fstat(source).st_size
+            except OSError:
+                pass
+            if detached and end == self._size:
+                # Nothing left the file, so this object still owns it.
+                self._finalizer = weakref.finalize(self, self._remove, self._path)
+            elif end < self._size:
+                e.add_note(
+                    f"Bytes [0, {end}) are still in {self._path}; bytes "
+                    f"[{end}, {self._size}) are at offset {start + end} of the output."
+                )
+            raise
+        return start
+
     def close(self) -> None:
         if self._finalizer is not None:
             self._finalizer()
@@ -131,10 +288,35 @@ def append(self, data: Union[CordBuffer, "Cord"]) -> None:
                 f"Can only append bytes, FileBackedData, or Cords, received {type(data)}"
             )
 
+    def file_backed_buffers(self) -> List[FileBackedData]:
+        """The file-backed buffers of the Cord, in order, with repeats."""
+        return [item for item in self._buffers if isinstance(item, FileBackedData)]
+
     def write_to_file(self, outfile: io.BufferedIOBase) -> None:
         """Write the Cord to a file."""
-        for item in self._buffers:
-            if isinstance(item, bytes):
-                outfile.write(item)
-            else:
-                item.write_to_file(outfile)
+        self._write_to_file(outfile)
+
+    def _write_to_file(
+        self, outfile: io.BufferedIOBase, move: Container[FileBackedData] = ()
+    ) -> None:
+        """Write the Cord to ``outfile``, moving the buffers in ``move`` into it when they
+        can be (see ``FileBackedData``). Each of them must appear once. If the write
+        fails, a note on the error lists the buffers already moved, whose only copy is now
+        in ``outfile``, and where they are."""
+        moved: List[str] = []
+        try:
+            for item in self._buffers:
+                if isinstance(item, bytes):
+                    outfile.write(item)
+                    continue
+                offset = item._move_to(outfile) if item in move else None
+                if offset is None:
+                    item.write_to_file(outfile)
+                else:
+                    moved.append(f"{item.path} ({len(item)} bytes) at offset {offset}")
+        except BaseException as e:
+            if moved:
+                e.add_note(
+                    "Moved into the output before the failure: " + "; ".join(moved)
+                )
+            raise