[BACKEND] File lock working directory to prevent multiple instances from polluting it (#227)

This commit is contained in:
Jesse Bannon 2022-09-08 23:41:09 -07:00 committed by GitHub
parent d0cbd57a7f
commit 5488501bed
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 83 additions and 7 deletions

View file

@ -1,5 +1,11 @@
import argparse import argparse
import errno
import fcntl
import os
import sys import sys
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import List from typing import List
from typing import Tuple from typing import Tuple
@ -8,6 +14,7 @@ from ytdl_sub.cli.main_args_parser import parser
from ytdl_sub.config.config_file import ConfigFile from ytdl_sub.config.config_file import ConfigFile
from ytdl_sub.config.preset import Preset from ytdl_sub.config.preset import Preset
from ytdl_sub.subscriptions.subscription import Subscription from ytdl_sub.subscriptions.subscription import Subscription
from ytdl_sub.utils.exceptions import ValidationException
from ytdl_sub.utils.file_handler import FileHandlerTransactionLog from ytdl_sub.utils.file_handler import FileHandlerTransactionLog
from ytdl_sub.utils.logger import Logger from ytdl_sub.utils.logger import Logger
@ -89,6 +96,41 @@ def _download_subscription_from_cli(
return subscription, subscription.download(dry_run=args.dry_run) return subscription, subscription.download(dry_run=args.dry_run)
@contextmanager
def _working_directory_lock(config: ConfigFile):
"""
Create and try to lock the file /tmp/working_directory_name
Raises
------
ValidationException
Lock is acquired from another process running ytdl-sub in the same working directory
OSError
Other lock error occurred
"""
working_directory_path = Path(os.getcwd()) / config.config_options.working_directory
lock_file_path = Path(tempfile.gettempdir()) / str(working_directory_path).replace("/", "_")
lock_file = open(lock_file_path, "w", encoding="utf-8")
try:
fcntl.lockf(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
except OSError as exc:
if exc.errno in (errno.EACCES, errno.EAGAIN):
raise ValidationException(
"Cannot run two instances of ytdl-sub "
"with the same working directory at the same time"
) from exc
lock_file.close()
raise exc
try:
yield
finally:
fcntl.flock(lock_file, fcntl.LOCK_UN)
lock_file.close()
def main() -> List[Tuple[Subscription, FileHandlerTransactionLog]]: def main() -> List[Tuple[Subscription, FileHandlerTransactionLog]]:
""" """
Entrypoint for ytdl-sub, without the error handling Entrypoint for ytdl-sub, without the error handling
@ -102,14 +144,16 @@ def main() -> List[Tuple[Subscription, FileHandlerTransactionLog]]:
config: ConfigFile = ConfigFile.from_file_path(args.config).initialize() config: ConfigFile = ConfigFile.from_file_path(args.config).initialize()
transaction_logs: List[Tuple[Subscription, FileHandlerTransactionLog]] = [] transaction_logs: List[Tuple[Subscription, FileHandlerTransactionLog]] = []
if args.subparser == "sub":
transaction_logs = _download_subscriptions_from_yaml_files(config=config, args=args)
# One-off download with _working_directory_lock(config=config):
if args.subparser == "dl": if args.subparser == "sub":
transaction_logs.append( transaction_logs = _download_subscriptions_from_yaml_files(config=config, args=args)
_download_subscription_from_cli(config=config, args=args, extra_args=extra_args)
) # One-off download
if args.subparser == "dl":
transaction_logs.append(
_download_subscription_from_cli(config=config, args=args, extra_args=extra_args)
)
for subscription, transaction_log in transaction_logs: for subscription, transaction_log in transaction_logs:
if transaction_log.is_empty: if transaction_log.is_empty:

View file

@ -0,0 +1,32 @@
import os
import time
import pytest
from ytdl_sub.cli.main import _working_directory_lock
from ytdl_sub.config.config_file import ConfigFile
from ytdl_sub.utils.exceptions import ValidationException
@pytest.fixture
def config() -> ConfigFile:
return ConfigFile(
name="config", value={"configuration": {"working_directory": "test"}, "presets": {}}
)
def test_working_directory_lock(config: ConfigFile):
new_pid = os.fork()
if new_pid == 0: # is child
with _working_directory_lock(config=config):
time.sleep(3)
return
time.sleep(1)
with pytest.raises(ValidationException, match="Cannot run two instances of ytdl-sub"):
with _working_directory_lock(config=config):
time.sleep(1)
time.sleep(3)
with _working_directory_lock(config=config):
pass