From 4e3996f3f8ddd3bc65ed6d36f93ed3465d1112ef Mon Sep 17 00:00:00 2001 From: Michael Hudson-Doyle Date: Tue, 5 May 2020 16:14:11 +1200 Subject: [PATCH] improve with_context helper a bit, use it more the decorated function must now be called with keyword arguments --- subiquity/controllers/cmdlist.py | 4 +- subiquity/controllers/filesystem.py | 86 ++++++++++++------------ subiquity/controllers/identity.py | 2 +- subiquity/controllers/installprogress.py | 36 +++++----- subiquity/controllers/network.py | 4 +- subiquity/controllers/proxy.py | 2 +- subiquity/controllers/reboot.py | 38 +++++------ subiquity/controllers/refresh.py | 2 +- subiquity/controllers/snaplist.py | 46 +++++++------ subiquity/controllers/ssh.py | 9 ++- subiquitycore/context.py | 32 +++++++-- subiquitycore/controllers/network.py | 4 +- 12 files changed, 145 insertions(+), 120 deletions(-) diff --git a/subiquity/controllers/cmdlist.py b/subiquity/controllers/cmdlist.py index f644f1d2..50621d46 100644 --- a/subiquity/controllers/cmdlist.py +++ b/subiquity/controllers/cmdlist.py @@ -38,7 +38,7 @@ class CmdListController(NoUIController): @with_context() async def run(self, context): for i, cmd in enumerate(self.cmds): - with self.context.child("command_{}".format(i), cmd): + with context.child("command_{}".format(i), cmd): if isinstance(cmd, str): cmd = ['sh', '-c', cmd] await arun_command( @@ -58,4 +58,4 @@ class LateController(CmdListController): @with_context() async def apply_autoinstall_config(self, context): - await self.run(context) + await self.run(context=context) diff --git a/subiquity/controllers/filesystem.py b/subiquity/controllers/filesystem.py index 96b62aec..8f6f38c6 100644 --- a/subiquity/controllers/filesystem.py +++ b/subiquity/controllers/filesystem.py @@ -104,11 +104,11 @@ class FilesystemController(SubiquityController): self.ai_data = data @with_context() - async def apply_autoinstall_config(self, context): + async def apply_autoinstall_config(self, context=None): self.stop_listening_udev() await self._start_task await self._probe_task.wait() - self.convert_autoinstall_config() + self.convert_autoinstall_config(context=context) if not self.model.is_root_mounted(): raise Exception("autoinstall config did not mount root") if self.model.needs_bootloader_partition(): @@ -116,7 +116,8 @@ class FilesystemController(SubiquityController): "autoinstall config did not create needed bootloader " "partition") - async def _probe_once(self, restricted): + @with_context(name='probe_once', description='restricted={restricted}') + async def _probe_once(self, *, context, restricted): if restricted: probe_types = {'blockdev'} fname = 'probe-data-restricted.json' @@ -133,52 +134,49 @@ class FilesystemController(SubiquityController): self.app.note_file_for_apport(key, fpath) self.model.load_probe_data(storage) - async def _probe(self): - with self.context.child("_probe") as context: - async with self.app.install_lock_file.shared(): - self._crash_reports = {} - if isinstance(self.ui.body, ProbingFailed): - self.ui.set_body(SlowProbing(self)) - schedule_task(self._wait_for_probing()) - for (restricted, kind) in [ - (False, ErrorReportKind.BLOCK_PROBE_FAIL), - (True, ErrorReportKind.DISK_PROBE_FAIL), - ]: - try: - desc = "restricted={}".format(restricted) - with context.child("probe_once", desc): - await self._probe_once_task.start(restricted) - # We wait on the task directly here, not - # self._probe_once_task.wait as if _probe_once_task - # gets cancelled, we should be cancelled too. - await asyncio.wait_for( - self._probe_once_task.task, 15.0) - except asyncio.CancelledError: - # asyncio.CancelledError is a subclass of Exception in - # Python 3.6 (sadface) - raise - except Exception: - block_discover_log.exception( - "block probing failed restricted=%s", restricted) - report = self.app.make_apport_report( - kind, "block probing", interrupt=False) - self._crash_reports[restricted] = report - continue - break + @with_context() + async def _probe(self, *, context=None): + async with self.app.install_lock_file.shared(): + self._crash_reports = {} + if isinstance(self.ui.body, ProbingFailed): + self.ui.set_body(SlowProbing(self)) + schedule_task(self._wait_for_probing()) + for (restricted, kind) in [ + (False, ErrorReportKind.BLOCK_PROBE_FAIL), + (True, ErrorReportKind.DISK_PROBE_FAIL), + ]: + try: + await self._probe_once_task.start( + context=context, restricted=restricted) + # We wait on the task directly here, not + # self._probe_once_task.wait as if _probe_once_task + # gets cancelled, we should be cancelled too. + await asyncio.wait_for(self._probe_once_task.task, 15.0) + except asyncio.CancelledError: + # asyncio.CancelledError is a subclass of Exception in + # Python 3.6 (sadface) + raise + except Exception: + block_discover_log.exception( + "block probing failed restricted=%s", restricted) + report = self.app.make_apport_report( + kind, "block probing", interrupt=False) + self._crash_reports[restricted] = report + continue + break - def convert_autoinstall_config(self): + @with_context() + def convert_autoinstall_config(self, context=None): log.debug("self.ai_data = %s", self.ai_data) if 'layout' in self.ai_data: layout = self.ai_data['layout'] - with self.context.child("applying_autoinstall"): - meth = getattr(self, "guided_" + layout['name']) - disk = self.model.disk_for_match( - self.model.all_disks(), - layout.get("match", {'size': 'largest'})) - meth(disk) + meth = getattr(self, "guided_" + layout['name']) + disk = self.model.disk_for_match( + self.model.all_disks(), + layout.get("match", {'size': 'largest'})) + meth(disk) elif 'config' in self.ai_data: - with self.context.child("applying_autoinstall"): - self.model.apply_autoinstall_config(self.ai_data['config']) + self.model.apply_autoinstall_config(self.ai_data['config']) self.model.grub = self.ai_data.get('grub', {}) self.model.swap = self.ai_data.get('swap') diff --git a/subiquity/controllers/identity.py b/subiquity/controllers/identity.py index 942ce7c9..176b3146 100644 --- a/subiquity/controllers/identity.py +++ b/subiquity/controllers/identity.py @@ -45,7 +45,7 @@ class IdentityController(SubiquityController): self.model.add_user(data) @with_context() - async def apply_autoinstall_config(self, context): + async def apply_autoinstall_config(self, context=None): if not self.model.user: if 'user-data' not in self.app.autoinstall_config: raise Exception("no identity data provided") diff --git a/subiquity/controllers/installprogress.py b/subiquity/controllers/installprogress.py index 954af1de..8f392037 100644 --- a/subiquity/controllers/installprogress.py +++ b/subiquity/controllers/installprogress.py @@ -241,7 +241,7 @@ class InstallProgressController(SubiquityController): return curtin_cmd @with_context(description="umounting /target dir") - async def unmount_target(self, context, target): + async def unmount_target(self, *, context, target): cmd = [ sys.executable, '-m', 'curtin', 'unmount', '-t', target, @@ -254,7 +254,7 @@ class InstallProgressController(SubiquityController): @with_context( description="installing system", level="INFO", childlevel="DEBUG") - async def curtin_install(self, context): + async def curtin_install(self, *, context): log.debug('curtin_install') self.install_state = InstallState.RUNNING self.curtin_event_contexts[''] = context @@ -288,7 +288,7 @@ class InstallProgressController(SubiquityController): pass @with_context() - async def install(self, context): + async def install(self, *, context): context.set('is-install-context', True) try: await asyncio.wait( @@ -297,16 +297,17 @@ class InstallProgressController(SubiquityController): await self.confirmation.wait() if os.path.exists(self.model.target): - await self.unmount_target(context, self.model.target) + await self.unmount_target( + context=context, target=self.model.target) - await self.curtin_install(context) + await self.curtin_install(context=context) await asyncio.wait( {e.wait() for e in self.model.postinstall_events}) - await self.drain_curtin_events(context) + await self.drain_curtin_events(context=context) - await self.postinstall(context) + await self.postinstall(context=context) self.ui.set_header(_("Installation complete!")) self.progress_view.set_status(_("Finished install!")) @@ -314,7 +315,7 @@ class InstallProgressController(SubiquityController): if self.model.network.has_network: self.progress_view.update_running() - await self.run_unattended_upgrades(context) + await self.run_unattended_upgrades(context=context) self.progress_view.update_done() except Exception: @@ -326,7 +327,7 @@ class InstallProgressController(SubiquityController): await self.install_task self.app.next_screen() - async def drain_curtin_events(self, context): + async def drain_curtin_events(self, *, context): waited = 0.0 while self.progress_view.ongoing and waited < 5.0: await asyncio.sleep(0.1) @@ -337,30 +338,29 @@ class InstallProgressController(SubiquityController): @with_context( description="final system configuration", level="INFO", childlevel="DEBUG") - async def postinstall(self, context): + async def postinstall(self, *, context): autoinstall_path = os.path.join( self.app.root, 'var/log/installer/autoinstall-user-data') autoinstall_config = "#cloud-config\n" + yaml.dump( {"autoinstall": self.app.make_autoinstall()}) write_file(autoinstall_path, autoinstall_config, mode=0o600) - await self.configure_cloud_init(context) + await self.configure_cloud_init(context=context) packages = [] if self.model.ssh.install_server: packages = ['openssh-server'] packages.extend(self.app.base_model.packages) for package in packages: - subcontext = context.child( - "install_{}".format(package), - "installing {}".format(package)) - with subcontext: - await self.install_package(package) - await self.restore_apt_config(context) + await self.install_package(context=context, package=package) + await self.restore_apt_config(context=context) @with_context(description="configuring cloud-init") async def configure_cloud_init(self, context): await run_in_thread(self.model.configure_cloud_init) - async def install_package(self, package): + @with_context( + name="install_{package}", + description="installing {package}") + async def install_package(self, *, context, package): if self.opts.dry_run: cmd = ["sleep", str(2/self.app.scale_factor)] else: diff --git a/subiquity/controllers/network.py b/subiquity/controllers/network.py index d7a990c2..77d599fc 100644 --- a/subiquity/controllers/network.py +++ b/subiquity/controllers/network.py @@ -155,9 +155,9 @@ class NetworkController(NetworkController, SubiquityController): self.model.has_network = bool( self.network_event_receiver.default_routes) - async def _apply_config(self, context=None, *, silent): + async def _apply_config(self, *, context=None, silent=False): try: - await super()._apply_config(context, silent=silent) + await super()._apply_config(context=context, silent=silent) except asyncio.CancelledError: # asyncio.CancelledError is a subclass of Exception in # Python 3.6 (sadface) diff --git a/subiquity/controllers/proxy.py b/subiquity/controllers/proxy.py index 8d22ae82..f4620fdf 100644 --- a/subiquity/controllers/proxy.py +++ b/subiquity/controllers/proxy.py @@ -43,7 +43,7 @@ class ProxyController(SubiquityController): self.signal.emit_signal('network-proxy-set') @with_context() - async def apply_autoinstall_config(self, context): + async def apply_autoinstall_config(self, context=None): # XXX want to wait until signal sent by .start() has been seen # by everything; don't have a way to do that today. pass diff --git a/subiquity/controllers/reboot.py b/subiquity/controllers/reboot.py index 8bf7d029..c7f5bc4e 100644 --- a/subiquity/controllers/reboot.py +++ b/subiquity/controllers/reboot.py @@ -36,25 +36,25 @@ class RebootController(SubiquityController): def interactive(self): return self.app.interactive() - async def copy_logs_to_target(self): - with self.context.child("copy_logs_to_target"): - if self.opts.dry_run and 'copy-logs-fail' in self.app.debug_flags: - raise PermissionError() - target_logs = os.path.join( - self.app.base_model.target, 'var/log/installer') - if self.opts.dry_run: - os.makedirs(target_logs, exist_ok=True) - else: + @with_context() + async def copy_logs_to_target(self, context): + if self.opts.dry_run and 'copy-logs-fail' in self.app.debug_flags: + raise PermissionError() + target_logs = os.path.join( + self.app.base_model.target, 'var/log/installer') + if self.opts.dry_run: + os.makedirs(target_logs, exist_ok=True) + else: + await arun_command( + ['cp', '-aT', '/var/log/installer', target_logs]) + journal_txt = os.path.join(target_logs, 'installer-journal.txt') + try: + with open(journal_txt, 'w') as output: await arun_command( - ['cp', '-aT', '/var/log/installer', target_logs]) - journal_txt = os.path.join(target_logs, 'installer-journal.txt') - try: - with open(journal_txt, 'w') as output: - await arun_command( - ['journalctl', '-b'], - stdout=output, stderr=subprocess.STDOUT) - except Exception: - log.exception("saving journal failed") + ['journalctl', '-b'], + stdout=output, stderr=subprocess.STDOUT) + except Exception: + log.exception("saving journal failed") def reboot(self): if self.opts.dry_run: @@ -66,7 +66,7 @@ class RebootController(SubiquityController): @with_context() async def apply_autoinstall_config(self, context): - await self.copy_logs_to_target() + await self.copy_logs_to_target(context=context) self.reboot() async def _run(self): diff --git a/subiquity/controllers/refresh.py b/subiquity/controllers/refresh.py index 9580c0e8..8713425e 100644 --- a/subiquity/controllers/refresh.py +++ b/subiquity/controllers/refresh.py @@ -97,7 +97,7 @@ class RefreshController(SubiquityController): return if self.check_state != CheckState.AVAILABLE: return - change_id = await self.start_update(context) + change_id = await self.start_update(context=context) while True: try: change = await self.get_progress(change_id) diff --git a/subiquity/controllers/snaplist.py b/subiquity/controllers/snaplist.py index 194fecd1..ae873b9c 100644 --- a/subiquity/controllers/snaplist.py +++ b/subiquity/controllers/snaplist.py @@ -20,6 +20,7 @@ import requests.exceptions from subiquitycore.async_helpers import ( schedule_task, ) +from subiquitycore.context import with_context from subiquitycore.controller import ( Skip, ) @@ -62,34 +63,34 @@ class SnapdSnapInfoLoader: while self.pending_snaps: snap = self.pending_snaps.pop(0) task = self.tasks[snap] = schedule_task( - self._fetch_info_for_snap(snap)) + self._fetch_info_for_snap(snap=snap)) await task - async def _load_list(self): - with self.context.child("list"): - try: - result = await self.snapd.get( - 'v2/find', section=self.store_section) - except requests.exceptions.RequestException: - log.exception("loading list of snaps failed") - self.failed = True - return - self.model.load_find_data(result) - self.snap_list_fetched = True + @with_context(name="list") + async def _load_list(self, context=None): + try: + result = await self.snapd.get( + 'v2/find', section=self.store_section) + except requests.exceptions.RequestException: + log.exception("loading list of snaps failed") + self.failed = True + return + self.model.load_find_data(result) + self.snap_list_fetched = True def stop(self): if self.main_task is not None: self.main_task.cancel() - async def _fetch_info_for_snap(self, snap): - with self.context.child("fetch").child(snap.name): - try: - data = await self.snapd.get('v2/find', name=snap.name) - except requests.exceptions.RequestException: - log.exception("loading snap info failed") - # XXX something better here? - return - self.model.load_info_data(data) + @with_context(name="fetch/{snap.name}") + async def _fetch_info_for_snap(self, snap, context=None): + try: + data = await self.snapd.get('v2/find', name=snap.name) + except requests.exceptions.RequestException: + log.exception("loading snap info failed") + # XXX something better here? + return + self.model.load_info_data(data) def get_snap_list_task(self): return self.tasks[None] @@ -98,7 +99,8 @@ class SnapdSnapInfoLoader: if snap not in self.tasks: if snap in self.pending_snaps: self.pending_snaps.remove(snap) - self.tasks[snap] = schedule_task(self._fetch_info_for_snap(snap)) + self.tasks[snap] = schedule_task( + self._fetch_info_for_snap(snap=snap)) return self.tasks[snap] diff --git a/subiquity/controllers/ssh.py b/subiquity/controllers/ssh.py index cfaa740b..470d332d 100644 --- a/subiquity/controllers/ssh.py +++ b/subiquity/controllers/ssh.py @@ -17,6 +17,7 @@ import logging import subprocess from subiquitycore.async_helpers import schedule_task +from subiquitycore.context import with_context from subiquitycore import utils from subiquity.controller import SubiquityController @@ -94,7 +95,10 @@ class SSHController(SubiquityController): raise subprocess.CalledProcessError(cp.returncode, cmd) return cp - async def _fetch_ssh_keys(self, user_spec): + @with_context( + name="ssh_import_id", + description="{user_spec[ssh_import_id]}:{user_spec[import_username]}") + async def _fetch_ssh_keys(self, *, context, user_spec): ssh_import_id = "{ssh_import_id}:{import_username}".format(**user_spec) with self.context.child("ssh_import_id", ssh_import_id): try: @@ -127,7 +131,8 @@ class SSHController(SubiquityController): user_spec, ssh_import_id, key_material, fingerprints) def fetch_ssh_keys(self, user_spec): - self._fetch_task = schedule_task(self._fetch_ssh_keys(user_spec)) + self._fetch_task = schedule_task( + self._fetch_ssh_keys(user_spec=user_spec)) def done(self, result): log.debug("SSHController.done next_screen result=%s", result) diff --git a/subiquitycore/context.py b/subiquitycore/context.py index 4a8f2b99..2743d751 100644 --- a/subiquitycore/context.py +++ b/subiquitycore/context.py @@ -15,6 +15,8 @@ import asyncio import enum +import functools +import inspect class Status(enum.Enum): @@ -117,12 +119,30 @@ def with_context(name=None, description="", **context_kw): if name is None: name = meth.__name__ - async def decorated(self, context=None, *args, **kw): + def convargs(self, kw): + context = kw.get('context') if context is None: context = self.context - manager = context.child( - name, description=description.format(**kw), **context_kw) - with manager as subcontext: - await meth(self, subcontext, *args, **kw) - return decorated + kw['context'] = context.child( + name=name.format(**kw), + description=description.format(**kw), + **context_kw) + return kw + + @functools.wraps(meth) + def decorated_sync(self, **kw): + kw = convargs(self, kw) + with kw['context']: + return meth(self, **kw) + + @functools.wraps(meth) + async def decorated_async(self, **kw): + kw = convargs(self, kw) + with kw['context']: + return await meth(self, **kw) + + if inspect.iscoroutinefunction(meth): + return decorated_async + else: + return decorated_sync return decorate diff --git a/subiquitycore/controllers/network.py b/subiquitycore/controllers/network.py index c11658b5..c540856c 100644 --- a/subiquitycore/controllers/network.py +++ b/subiquitycore/controllers/network.py @@ -317,7 +317,7 @@ class NetworkController(BaseController): return os.path.join(self.root, 'etc/netplan', netplan_config_file_name) def apply_config(self, context=None, silent=False): - self.apply_config_task.start_sync(context, silent=silent) + self.apply_config_task.start_sync(context=context, silent=silent) async def _down_devs(self, devs): for dev in devs: @@ -358,7 +358,7 @@ class NetworkController(BaseController): @with_context( name="apply_config", description="silent={silent}", level="INFO") - async def _apply_config(self, context, *, silent): + async def _apply_config(self, *, context, silent): devs_to_delete = [] devs_to_down = [] dhcp_device_versions = []