From 26651626d7c454ec1027cbd044570b72de6106f0 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Thu, 12 Jan 2023 13:10:57 +0545
Subject: [PATCH] ProcessInfo: save atomically (#104)

---
 src/dvc_task/proc/process.py | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/dvc_task/proc/process.py b/src/dvc_task/proc/process.py
index dd36d05..7564cf6 100644
--- a/src/dvc_task/proc/process.py
+++ b/src/dvc_task/proc/process.py
@@ -5,6 +5,7 @@
 import os
 import shlex
 import subprocess  # nosec B404
+import tempfile
 from contextlib import AbstractContextManager, ExitStack
 from dataclasses import asdict, dataclass
 from typing import Any, Dict, List, Optional, Union
@@ -12,7 +13,7 @@
 from funcy import cached_property
 from shortuuid import uuid
 
-from ..contrib.kombu_filesystem import LOCK_EX, LOCK_SH, lock, unlock
+from ..contrib.kombu_filesystem import LOCK_SH, lock, unlock
 from ..utils import makedirs
 from .exceptions import TimeoutExpired
 
@@ -50,12 +51,17 @@ def asdict(self) -> Dict[str, Any]:
 
     def dump(self, filename: str) -> None:
         """Dump the process information into a file."""
-        with open(filename, "w", encoding="utf-8") as fobj:
-            lock(fobj, LOCK_EX)
-            try:
-                json.dump(self.asdict(), fobj)
-            finally:
-                unlock(fobj)
+        directory, file = os.path.split(filename)
+        with tempfile.NamedTemporaryFile(
+            mode="w",
+            encoding="utf-8",
+            dir=directory,
+            prefix=f"{file}.",
+            suffix=".tmp",
+            delete=False,
+        ) as tmp:
+            json.dump(self.asdict(), tmp)
+        os.replace(tmp.name, filename)
 
 
 class ManagedProcess(AbstractContextManager):