fix: cleanup edge cases and deduplicate subvol lookup

- fstab.py: clean up temp file on write/fsync OSError in _atomic_write
- atomic-gc: cache current_subvol once in rm command (match
garbage_collect)
- atomic-rebuild-uki: only rmdir if umount succeeded in cleanup_rebuild
- tests: add fsync failure cleanup test and rebuild-uki cleanup
behavioral test

Author: linux_name2

bin/atomic-gc

@@ -81,6 +81,7 @@ case "$COMMAND" in
         acquire_lock
         [[ ${#ARGS[@]} -eq 0 ]] && { echo "ERROR: Specify generation(s) to delete" >&2; exit 1; }
         ensure_btrfs_mounted || exit 1
+        current_subvol=$(get_current_subvol)
         # Validate all targets exist before deleting any
         for gen_id in "${ARGS[@]}"; do
             if [[ ! -f "${ESP}/EFI/Linux/arch-${gen_id}.efi" ]] && \
@@ -98,7 +99,7 @@ case "$COMMAND" in
             warn_orphan_homes "${ARGS[@]}"
         fi
         for gen_id in "${ARGS[@]}"; do
-            delete_generation "$gen_id" "$DRY_RUN"
+            delete_generation "$gen_id" "$DRY_RUN" "$current_subvol"
         done
         ;;
     gc)

bin/atomic-rebuild-uki

@@ -107,8 +107,9 @@ acquire_lock
 cleanup_rebuild() {
     set +e
     [[ -n "${MOUNT_DIR:-}" ]] && {
-        mountpoint -q "$MOUNT_DIR" 2>/dev/null && umount "$MOUNT_DIR"
-        rmdir "$MOUNT_DIR" 2>/dev/null
+        if mountpoint -q "$MOUNT_DIR" 2>/dev/null && umount "$MOUNT_DIR"; then
+            rmdir "$MOUNT_DIR" 2>/dev/null || true
+        fi
     }
     mountpoint -q "$BTRFS_MOUNT" 2>/dev/null && umount "$BTRFS_MOUNT"
     [[ -n "${LOCK_FD:-}" ]] && exec {LOCK_FD}>&-

lib/atomic/fstab.py

@@ -171,8 +171,13 @@ def _atomic_write(path: Path, entries: list) -> None:
     try:
         os.write(fd, content)
         os.fsync(fd)
-    finally:
+    except OSError:
         os.close(fd)
+        try:
+            tmp.unlink(missing_ok=True)
+        except OSError:
+            pass
+        raise
 
     try:
         os.chown(tmp, original_stat.st_uid, original_stat.st_gid)

tests/test_fstab.py

@@ -474,3 +474,25 @@ def fail_chmod_on_tmp(file_path, mode, *args, **kwargs):
         # Verify that the temp file was cleaned up despite the error
         leftovers = list(tmp_path.glob(".fstab.*.tmp"))
         assert leftovers == [], f"Temp files left behind: {leftovers}"
+
+    def test_atomic_write_fsync_failure_cleanup(self, tmp_path, monkeypatch):
+        """Simulate fsync failure to verify temp file cleanup and OSError propagation."""
+        path = tmp_path / "fstab"
+        path.write_text("UUID=x / btrfs rw,subvol=old 0 0\n")
+
+        original_fsync = os.fsync
+
+        def fail_fsync(fd):
+            raise OSError("Simulated fsync failure")
+
+        monkeypatch.setattr(os, "fsync", fail_fsync)
+
+        with pytest.raises(OSError, match="Simulated fsync failure"):
+            update_fstab(str(path), "old", "new")
+
+        # Verify that the temp file was cleaned up despite the error
+        leftovers = list(tmp_path.glob(".fstab.*.tmp"))
+        assert leftovers == [], f"Temp files left behind: {leftovers}"
+
+        # Verify original fstab is unchanged
+        assert path.read_text() == "UUID=x / btrfs rw,subvol=old 0 0\n"

tests/test_rebuild_uki.sh

@@ -134,6 +134,70 @@ grep -q 'rmdir.*MOUNT_DIR' "$SCRIPT" && ok "cleanup removes MOUNT_DIR" || fail "
 grep -q 'umount.*BTRFS_MOUNT' "$SCRIPT" && ok "cleanup unmounts BTRFS" || fail "cleanup unmounts BTRFS"
 assert_contains "cleanup closes lock" 'LOCK_FD' "$_script_content"
 
+# ── Cleanup trap: behavioral test — rmdir only if umount succeeds ──
+
+section "Cleanup trap: behavioral test"
+
+_RMDIR_LOG="${TESTDIR}/rmdir_calls.log"
+_UMOUNT_LOG="${TESTDIR}/umount_calls.log"
+
+# Test 1: umount fails → rmdir must NOT be called
+make_mock umount "echo \"\$*\" >> '${_UMOUNT_LOG}'; exit 1"
+make_mock rmdir "echo \"\$*\" >> '${_RMDIR_LOG}'; exit 0"
+make_mock mountpoint "exit 0"
+
+_TEST_SCRIPT="${TESTDIR}/test-cleanup"
+cat > "$_TEST_SCRIPT" << 'EOF'
+#!/bin/bash
+set -euo pipefail
+MOUNT_DIR="/tmp/fake-mount-dir"
+BTRFS_MOUNT="/run/atomic/temp_btrfs"
+LOCK_FD=""
+EOF
+
+# Append the actual cleanup function from the real script
+sed -n '/^cleanup_rebuild()/,/^}/p' "$SCRIPT" >> "$_TEST_SCRIPT"
+
+cat >> "$_TEST_SCRIPT" << 'EOF'
+trap cleanup_rebuild EXIT
+exit 0
+EOF
+chmod +x "$_TEST_SCRIPT"
+
+PATH="${MOCK_BIN}:${PATH}" bash "$_TEST_SCRIPT" 2>/dev/null || true
+
+if [[ -f "$_UMOUNT_LOG" ]]; then
+    ok "cleanup calls umount (fail path)"
+else
+    fail "cleanup calls umount (fail path)"
+fi
+
+if [[ ! -f "$_RMDIR_LOG" ]]; then
+    ok "rmdir NOT called when umount fails"
+else
+    fail "rmdir NOT called when umount fails (called with: $(cat "$_RMDIR_LOG"))"
+fi
+
+# Test 2: umount succeeds → rmdir IS called
+rm -f "$_RMDIR_LOG" "$_UMOUNT_LOG"
+
+make_mock umount "echo \"\$*\" >> '${_UMOUNT_LOG}'; exit 0"
+make_mock rmdir "echo \"\$*\" >> '${_RMDIR_LOG}'; exit 0"
+
+PATH="${MOCK_BIN}:${PATH}" bash "$_TEST_SCRIPT" 2>/dev/null || true
+
+if [[ -f "$_UMOUNT_LOG" ]]; then
+    ok "cleanup calls umount (success path)"
+else
+    fail "cleanup calls umount (success path)"
+fi
+
+if [[ -f "$_RMDIR_LOG" ]]; then
+    ok "rmdir called when umount succeeds"
+else
+    fail "rmdir called when umount succeeds"
+fi
+
 # ── UKI path construction: verify in real script ──
 
 section "UKI path construction: verify in real script"