Skip to content

Commit

Permalink
ProcessInfo: save atomically (#104)
Browse files Browse the repository at this point in the history
  • Loading branch information
skshetry authored Jan 12, 2023
1 parent 59caee6 commit 2665162
Showing 1 changed file with 13 additions and 7 deletions.
20 changes: 13 additions & 7 deletions src/dvc_task/proc/process.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,15 @@
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

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

Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 2665162

Please sign in to comment.