Use context manager for locking clvm files

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)

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:
     """