From f6743197fee1cbd711f03fc33af198a1db882f2f Mon Sep 17 00:00:00 2001
From: Kyle Altendorf <sda@fstab.net>
Date: Sat, 7 May 2022 10:26:01 -0400
Subject: [PATCH] Use context manager for locking clvm files

---
 chia/util/lock.py                | 11 +++++------
 chia/wallet/puzzles/load_clvm.py |  6 ++----
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/chia/util/lock.py b/chia/util/lock.py
index dec44b32f995..8bed53cb341a 100644
--- a/chia/util/lock.py
+++ b/chia/util/lock.py
@@ -1,6 +1,7 @@
+import contextlib
 import os
 import time
-from typing import Callable, Optional, TextIO, TypeVar
+from typing import Iterator, Optional, TextIO, TypeVar
 
 T = TypeVar("T")
 
@@ -22,15 +23,13 @@ def create_exclusive_lock(lockfile: str) -> Optional[TextIO]:
     return f
 
 
-def with_lock(lock_filename: str, run: Callable[[], T]) -> T:
+@contextlib.contextmanager
+def lock_by_path(lock_filename: str) -> Iterator[None]:
     """
     Ensure that this process and this thread is the only one operating on the
     resource associated with lock_filename systemwide.
-
-    Pass through the result of run after exiting the lock.
     """
 
-    lock_file = None
     while True:
         lock_file = create_exclusive_lock(lock_filename)
         if lock_file is not None:
@@ -39,7 +38,7 @@ def with_lock(lock_filename: str, run: Callable[[], T]) -> T:
         time.sleep(0.1)
 
     try:
-        return run()
+        yield
     finally:
         lock_file.close()
         os.remove(lock_filename)
diff --git a/chia/wallet/puzzles/load_clvm.py b/chia/wallet/puzzles/load_clvm.py
index 89429a9f8b6f..89e2a2a5fcde 100644
--- a/chia/wallet/puzzles/load_clvm.py
+++ b/chia/wallet/puzzles/load_clvm.py
@@ -6,7 +6,7 @@
 
 import pkg_resources
 from chia.types.blockchain_format.program import Program, SerializedProgram
-from chia.util.lock import with_lock
+from chia.util.lock import lock_by_path
 from clvm_tools_rs import compile_clvm as compile_clvm_rust
 
 
@@ -67,11 +67,9 @@ def sha256file(f):
 
 
 def compile_clvm(full_path, output, search_paths=[]):
-    def do_compile():
+    with lock_by_path(f"{full_path}.lock"):
         compile_clvm_in_lock(full_path, output, search_paths)
 
-    with_lock(f"{full_path}.lock", do_compile)
-
 
 def load_serialized_clvm(clvm_filename, package_or_requirement=__name__) -> SerializedProgram:
     """