From e06409fdceedacd5e457d75feb2b490ec253dccc Mon Sep 17 00:00:00 2001
From: Daniel Langbein <daniel@systemli.org>
Date: Sun, 19 Mar 2023 20:10:06 +0100
Subject: [PATCH] refactor

---
 .../active_receive/btrfs_receive_chunks_active.py   |  3 ++-
 .../active_receive/exec_receive_chunks_active.py    | 13 +++++++++++++
 .../active_receive/exec_send_chunks_passive.py      | 13 +++++++++++++
 .../active_send/exec_receive_chunks_passive.py      | 13 +++++++++++++
 .../active_send/exec_send_chunks_active.py          | 13 +++++++++++++
 5 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/src/subprocess_util/chunked_transfer/active_receive/btrfs_receive_chunks_active.py b/src/subprocess_util/chunked_transfer/active_receive/btrfs_receive_chunks_active.py
index e0e893a..f482888 100644
--- a/src/subprocess_util/chunked_transfer/active_receive/btrfs_receive_chunks_active.py
+++ b/src/subprocess_util/chunked_transfer/active_receive/btrfs_receive_chunks_active.py
@@ -26,7 +26,8 @@ def main():
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(prog='btrfs-receive-chunks-active')
+    parser = argparse.ArgumentParser(prog='btrfs-receive-chunks-active',
+                                     description='Either CHILD_SUBVOLUME or SOURCE_CHUNK_DIR must be given')
 
     parser.add_argument('--chunk-dir',
                         help='Chunks are saved in this directory. '
diff --git a/src/subprocess_util/chunked_transfer/active_receive/exec_receive_chunks_active.py b/src/subprocess_util/chunked_transfer/active_receive/exec_receive_chunks_active.py
index 0bc80c2..b82b1c5 100644
--- a/src/subprocess_util/chunked_transfer/active_receive/exec_receive_chunks_active.py
+++ b/src/subprocess_util/chunked_transfer/active_receive/exec_receive_chunks_active.py
@@ -72,8 +72,21 @@ def exec_receive_chunks_active(
                 break
             chunk_no += 1
 
+    # Create the local chunk directory.
+    if target_chunk_dir.exists():
+        created_chunk_dir = False
+    else:
+        target_chunk_dir.mkdir(parents=True, exist_ok=False)
+        created_chunk_dir = True
+
     returncode = execute_consume_chunks(
         command=command,
         handle_chunks=handle_chunks,
     )
+
+    # If no errors occurred, delete the chunk directory.
+    # But only if we created it ourselves!
+    if returncode == 0 and created_chunk_dir:
+        target_chunk_dir.unlink(missing_ok=False)
+
     return returncode
diff --git a/src/subprocess_util/chunked_transfer/active_receive/exec_send_chunks_passive.py b/src/subprocess_util/chunked_transfer/active_receive/exec_send_chunks_passive.py
index 31478ed..610ee10 100644
--- a/src/subprocess_util/chunked_transfer/active_receive/exec_send_chunks_passive.py
+++ b/src/subprocess_util/chunked_transfer/active_receive/exec_send_chunks_passive.py
@@ -45,10 +45,23 @@ def exec_send_chunks_passive(
         chunk_transferred_sock.close()
         chunk_transferred_sock_path.unlink(missing_ok=False)
 
+    # Create the local chunk directory.
+    if source_chunk_dir.exists():
+        created_chunk_dir = False
+    else:
+        source_chunk_dir.mkdir(parents=True, exist_ok=False)
+        created_chunk_dir = True
+
     returncode = execute_produce_chunks(
         command=command,
         get_chunk_path=get_source_chunk_path,
         handle_chunk=handle_chunk,
         chunk_size=chunk_size,
     )
+
+    # If no errors occurred, delete the chunk directory.
+    # But only if we created it ourselves!
+    if returncode == 0 and created_chunk_dir:
+        source_chunk_dir.unlink(missing_ok=False)
+
     return returncode
diff --git a/src/subprocess_util/chunked_transfer/active_send/exec_receive_chunks_passive.py b/src/subprocess_util/chunked_transfer/active_send/exec_receive_chunks_passive.py
index 0f6481b..8ab39d7 100644
--- a/src/subprocess_util/chunked_transfer/active_send/exec_receive_chunks_passive.py
+++ b/src/subprocess_util/chunked_transfer/active_send/exec_receive_chunks_passive.py
@@ -40,8 +40,21 @@ def exec_receive_chunks_passive(
         sock.close()
         target_socket.unlink(missing_ok=False)
 
+    # Create the local chunk directory.
+    if target_chunk_dir.exists():
+        created_chunk_dir = False
+    else:
+        target_chunk_dir.mkdir(parents=True, exist_ok=False)
+        created_chunk_dir = True
+
     returncode = execute_consume_chunks(
         command=command,
         handle_chunks=handle_chunks,
     )
+
+    # If no errors occurred, delete the chunk directory.
+    # But only if we created it ourselves!
+    if returncode == 0 and created_chunk_dir:
+        target_chunk_dir.unlink(missing_ok=False)
+
     return returncode
diff --git a/src/subprocess_util/chunked_transfer/active_send/exec_send_chunks_active.py b/src/subprocess_util/chunked_transfer/active_send/exec_send_chunks_active.py
index 5f91f9f..8f08c50 100644
--- a/src/subprocess_util/chunked_transfer/active_send/exec_send_chunks_active.py
+++ b/src/subprocess_util/chunked_transfer/active_send/exec_send_chunks_active.py
@@ -41,10 +41,23 @@ def exec_send_chunks_active(
         source_chunk_path.unlink(missing_ok=False)
         repeat_until_successful(inform_cmd, usr_confirmation_socket)
 
+    # Create the local chunk directory.
+    if source_chunk_dir.exists():
+        created_chunk_dir = False
+    else:
+        source_chunk_dir.mkdir(parents=True, exist_ok=False)
+        created_chunk_dir = True
+
     returncode = execute_produce_chunks(
         command=command,
         get_chunk_path=get_source_chunk_path,
         handle_chunk=handle_chunk,
         chunk_size=chunk_size,
     )
+
+    # If no errors occurred, delete the chunk directory.
+    # But only if we created it ourselves!
+    if returncode == 0 and created_chunk_dir:
+        source_chunk_dir.unlink(missing_ok=False)
+
     return returncode