From cbb0aae3198cfe983346ed676bba7510c8172311 Mon Sep 17 00:00:00 2001
From: Marius Vollmer <marius.vollmer@gmail.com>
Date: Wed, 28 Aug 2024 16:34:48 +0300
Subject: [PATCH] WIP - conext cleanup

no contextlib.chdir on older OSes.
---
 pkg/storaged/btrfs/btrfs-tool.py | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/pkg/storaged/btrfs/btrfs-tool.py b/pkg/storaged/btrfs/btrfs-tool.py
index dd11d9d11a6f..88c19fb8e45b 100755
--- a/pkg/storaged/btrfs/btrfs-tool.py
+++ b/pkg/storaged/btrfs/btrfs-tool.py
@@ -30,11 +30,10 @@
 def debug(msg):
     subprocess.check_call(["logger", msg])
 
-def ensure_tmp_mp_dir():
-    os.makedirs(TMP_MP_DIR, mode=0o700, exist_ok=True)
-
 @contextlib.contextmanager
-def atomic_file(path):
+def mount_database():
+    path = TMP_MP_DIR + "/db"
+    os.makedirs(TMP_MP_DIR, mode=0o700, exist_ok=True)
     fd = os.open(path, os.O_RDWR | os.O_CREAT)
     fcntl.flock(fd, fcntl.LOCK_EX)
     data = os.read(fd, 100000)
@@ -48,6 +47,15 @@ def atomic_file(path):
     finally:
         os.close(fd)
 
+@contextlib.contextmanager
+def context_chdir(path):
+    old_cwd = os.getcwd()
+    os.chdir(path)
+    try:
+        yield
+    finally:
+        os.chdir(old_cwd)
+
 def list_filesystems():
     output = json.loads(subprocess.check_output(["lsblk", "-Jplno", "NAME,FSTYPE,UUID,MOUNTPOINTS"]))
     filesystems = {}
@@ -102,8 +110,7 @@ def remove_tmp_mountpoint(db, uuid):
             db[uuid] -= 1
 
 def remove_all_tmp_mountpoints():
-    ensure_tmp_mp_dir()
-    with atomic_file(TMP_MP_DIR + "/db") as db:
+    with mount_database() as db:
         for mp in set(tmp_mountpoints):
             remove_tmp_mountpoint(db, mp)
 
@@ -156,8 +163,7 @@ def get_usages(uuid):
 
 def poll(opt_mount, opt_repair):
     debug(f"POLL mount {opt_mount} repair {opt_repair}")
-    ensure_tmp_mp_dir()
-    with atomic_file(TMP_MP_DIR + "/db") as db:
+    with mount_database() as db:
         filesystems = list_filesystems()
         info = { }
         for fs in filesystems.values():
@@ -193,13 +199,12 @@ def cmd_poll(opt_mount):
 
 def cmd_do(uuid, cmd):
     debug(f"DO {uuid} {cmd}")
-    ensure_tmp_mp_dir()
-    with atomic_file(TMP_MP_DIR + "/db") as db:
+    with mount_database() as db:
         filesystems = list_filesystems()
         for fs in filesystems.values():
             if fs['uuid'] == uuid:
                 mp = force_mount_point(db, fs, True)
-                with contextlib.chdir(mp):
+                with context_chdir(mp):
                     subprocess.check_call(cmd)
 
 def cmd(args):