kookie
/
nomicon
1
0
Fork 0
Code Issues Projects Releases 1 Activity

maintainers/scripts/update.nix: switch to asyncio

Browse Source 
This will make it cleaner and also better respect SIGTERM.
wip/yesman
Jan Tojnar 4 years ago
parent 17f89667b3
commit 01b9d5371c
No known key found for this signature in database
GPG Key ID: 7FAB2A15F7A607A4
1 changed files with 113 additions and 67 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 180
      maintainers/scripts/update.py

180
maintainers/scripts/update.py
Unescape View File

@ -1,82 +1,142 @@
from typing import Dict, Generator, Tuple, Union
from __future__ import annotations
from typing import Dict, Generator, List, Optional, Tuple
import argparse
import asyncio
import contextlib
import concurrent.futures
import json
import os
import subprocess
import sys
import tempfile
import threading
updates: Dict[concurrent.futures.Future, Dict] = {}
TempDirs = Dict[str, Tuple[str, str, threading.Lock]]
thread_name_prefix = 'UpdateScriptThread'
class CalledProcessError(Exception):
process: asyncio.subprocess.Process
def eprint(*args, **kwargs):
print(*args, file=sys.stderr, **kwargs)
def run_update_script(package: Dict, commit: bool, temp_dirs: TempDirs) -> subprocess.CompletedProcess:
worktree: Union[None, str] = None
async def check_subprocess(*args, **kwargs):
"""
Emulate check argument of subprocess.run function.
"""
process = await asyncio.create_subprocess_exec(*args, **kwargs)
returncode = await process.wait()
if commit and 'commit' in package['supportedFeatures']:
thread_name = threading.current_thread().name
worktree, _branch, lock = temp_dirs[thread_name]
lock.acquire()
package['thread'] = thread_name
if returncode != 0:
error = CalledProcessError()
error.process = process
eprint(f" - {package['name']}: UPDATING ...")
raise error
return subprocess.run(package['updateScript'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True, cwd=worktree)
return process
@contextlib.contextmanager
def make_worktree() -> Generator[Tuple[str, str], None, None]:
with tempfile.TemporaryDirectory() as wt:
branch_name = f'update-{os.path.basename(wt)}'
target_directory = f'{wt}/nixpkgs'
async def run_update_script(merge_lock: asyncio.Lock, temp_dir: Optional[Tuple[str, str]], package: Dict, keep_going: bool):
worktree: Optional[str] = None
subprocess.run(['git', 'worktree', 'add', '-b', branch_name, target_directory], check=True)
yield (target_directory, branch_name)
subprocess.run(['git', 'worktree', 'remove', target_directory], check=True)
subprocess.run(['git', 'branch', '-D', branch_name], check=True)
if temp_dir is not None:
worktree, _branch = temp_dir
def commit_changes(worktree: str, branch: str, execution: subprocess.CompletedProcess) -> None:
changes = json.loads(execution.stdout)
for change in changes:
subprocess.run(['git', 'add'] + change['files'], check=True, cwd=worktree)
commit_message = '{attrPath}: {oldVersion} → {newVersion}'.format(**change)
subprocess.run(['git', 'commit', '-m', commit_message], check=True, cwd=worktree)
subprocess.run(['git', 'cherry-pick', branch], check=True)
eprint(f" - {package['name']}: UPDATING ...")
def merge_changes(package: Dict, future: concurrent.futures.Future, commit: bool, keep_going: bool, temp_dirs: TempDirs) -> None:
try:
execution = future.result()
if commit and 'commit' in package['supportedFeatures']:
thread_name = package['thread']
worktree, branch, lock = temp_dirs[thread_name]
commit_changes(worktree, branch, execution)
eprint(f" - {package['name']}: DONE.")
except subprocess.CalledProcessError as e:
update_process = await check_subprocess(*package['updateScript'], stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=worktree)
update_info = await update_process.stdout.read()
await merge_changes(merge_lock, package, update_info, temp_dir)
except KeyboardInterrupt as e:
eprint('Cancelling…')
raise asyncio.exceptions.CancelledError()
except CalledProcessError as e:
eprint(f" - {package['name']}: ERROR")
eprint()
eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
eprint()
eprint(e.stdout.decode('utf-8'))
stderr = await e.process.stderr.read()
eprint(stderr.decode('utf-8'))
with open(f"{package['pname']}.log", 'wb') as logfile:
logfile.write(e.stdout)
logfile.write(stderr)
eprint()
eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
if not keep_going:
sys.exit(1)
finally:
if commit and 'commit' in package['supportedFeatures']:
lock.release()
raise asyncio.exceptions.CancelledError()
def main(max_workers: int, keep_going: bool, commit: bool, packages: Dict) -> None:
with open(sys.argv[1]) as f:
@contextlib.contextmanager
def make_worktree() -> Generator[Tuple[str, str], None, None]:
with tempfile.TemporaryDirectory() as wt:
branch_name = f'update-{os.path.basename(wt)}'
target_directory = f'{wt}/nixpkgs'
subprocess.run(['git', 'worktree', 'add', '-b', branch_name, target_directory])
yield (target_directory, branch_name)
subprocess.run(['git', 'worktree', 'remove', '--force', target_directory])
subprocess.run(['git', 'branch', '-D', branch_name])
async def commit_changes(merge_lock: asyncio.Lock, worktree: str, branch: str, update_info: str) -> None:
changes = json.loads(update_info)
for change in changes:
# Git can only handle a single index operation at a time
async with merge_lock:
await check_subprocess('git', 'add', *change['files'], cwd=worktree)
commit_message = '{attrPath}: {oldVersion} → {newVersion}'.format(**change)
await check_subprocess('git', 'commit', '--quiet', '-m', commit_message, cwd=worktree)
await check_subprocess('git', 'cherry-pick', branch)
async def merge_changes(merge_lock: asyncio.Lock, package: Dict, update_info: str, temp_dir: Optional[Tuple[str, str]]) -> None:
if temp_dir is not None:
worktree, branch = temp_dir
await commit_changes(merge_lock, worktree, branch, update_info)
eprint(f" - {package['name']}: DONE.")
async def updater(temp_dir: Optional[Tuple[str, str]], merge_lock: asyncio.Lock, packages_to_update: asyncio.Queue[Optional[Dict]], keep_going: bool, commit: bool):
while True:
package = await packages_to_update.get()
if package is None:
# A sentinel received, we are done.
return
if not ('commit' in package['supportedFeatures']):
temp_dir = None
await run_update_script(merge_lock, temp_dir, package, keep_going)
async def start_updates(max_workers: int, keep_going: bool, commit: bool, packages: List[Dict]):
merge_lock = asyncio.Lock()
packages_to_update: asyncio.Queue[Optional[Dict]] = asyncio.Queue()
with contextlib.ExitStack() as stack:
temp_dirs: List[Optional[Tuple[str, str]]] = []
# Do not create more workers than there are packages.
num_workers = min(max_workers, len(packages))
# Set up temporary directories when using auto-commit.
for i in range(num_workers):
temp_dir = stack.enter_context(make_worktree()) if commit else None
temp_dirs.append(temp_dir)
# Fill up an update queue,
for package in packages:
await packages_to_update.put(package)
# Add sentinels, one for each worker.
# A workers will terminate when it gets sentinel from the queue.
for i in range(num_workers):
await packages_to_update.put(None)
# Prepare updater workers for each temp_dir directory.
# At most `num_workers` instances of `run_update_script` will be running at one time.
updaters = asyncio.gather(*[updater(temp_dir, merge_lock, packages_to_update, keep_going, commit) for temp_dir in temp_dirs])
try:
# Start updater workers.
await updaters
except asyncio.exceptions.CancelledError as e:
# When one worker is cancelled, cancel the others too.
updaters.cancel()
def main(max_workers: int, keep_going: bool, commit: bool, packages_path: str) -> None:
with open(packages_path) as f:
packages = json.load(f)
eprint()
@ -90,19 +150,7 @@ def main(max_workers: int, keep_going: bool, commit: bool, packages: Dict) -> No
eprint()
eprint('Running update for:')
with contextlib.ExitStack() as stack, concurrent.futures.ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix=thread_name_prefix) as executor:
if commit:
temp_dirs = {f'{thread_name_prefix}_{str(i)}': (*stack.enter_context(make_worktree()), threading.Lock()) for i in range(max_workers)}
else:
temp_dirs = {}
for package in packages:
updates[executor.submit(run_update_script, package, commit, temp_dirs)] = package
for future in concurrent.futures.as_completed(updates):
package = updates[future]
merge_changes(package, future, commit, keep_going, temp_dirs)
asyncio.run(start_updates(max_workers, keep_going, commit, packages))
eprint()
eprint('Packages updated!')
@ -122,8 +170,6 @@ if __name__ == '__main__':
try:
main(args.max_workers, args.keep_going, args.commit, args.packages)
except (KeyboardInterrupt, SystemExit) as e:
for update in updates:
update.cancel()
sys.exit(e.code if isinstance(e, SystemExit) else 130)
except KeyboardInterrupt as e:
# Let’s cancel outside of the main loop too.
sys.exit(130)

Write Preview
Loading…
Cancel
Save