improve with_context helper a bit, use it more

the decorated function must now be called with keyword arguments
This commit is contained in:
Michael Hudson-Doyle 2020-05-05 16:14:11 +12:00
parent 449f459ecb
commit 4e3996f3f8
12 changed files with 145 additions and 120 deletions

View File

@ -38,7 +38,7 @@ class CmdListController(NoUIController):
@with_context() @with_context()
async def run(self, context): async def run(self, context):
for i, cmd in enumerate(self.cmds): 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): if isinstance(cmd, str):
cmd = ['sh', '-c', cmd] cmd = ['sh', '-c', cmd]
await arun_command( await arun_command(
@ -58,4 +58,4 @@ class LateController(CmdListController):
@with_context() @with_context()
async def apply_autoinstall_config(self, context): async def apply_autoinstall_config(self, context):
await self.run(context) await self.run(context=context)

View File

@ -104,11 +104,11 @@ class FilesystemController(SubiquityController):
self.ai_data = data self.ai_data = data
@with_context() @with_context()
async def apply_autoinstall_config(self, context): async def apply_autoinstall_config(self, context=None):
self.stop_listening_udev() self.stop_listening_udev()
await self._start_task await self._start_task
await self._probe_task.wait() await self._probe_task.wait()
self.convert_autoinstall_config() self.convert_autoinstall_config(context=context)
if not self.model.is_root_mounted(): if not self.model.is_root_mounted():
raise Exception("autoinstall config did not mount root") raise Exception("autoinstall config did not mount root")
if self.model.needs_bootloader_partition(): if self.model.needs_bootloader_partition():
@ -116,7 +116,8 @@ class FilesystemController(SubiquityController):
"autoinstall config did not create needed bootloader " "autoinstall config did not create needed bootloader "
"partition") "partition")
async def _probe_once(self, restricted): @with_context(name='probe_once', description='restricted={restricted}')
async def _probe_once(self, *, context, restricted):
if restricted: if restricted:
probe_types = {'blockdev'} probe_types = {'blockdev'}
fname = 'probe-data-restricted.json' fname = 'probe-data-restricted.json'
@ -133,52 +134,49 @@ class FilesystemController(SubiquityController):
self.app.note_file_for_apport(key, fpath) self.app.note_file_for_apport(key, fpath)
self.model.load_probe_data(storage) self.model.load_probe_data(storage)
async def _probe(self): @with_context()
with self.context.child("_probe") as context: async def _probe(self, *, context=None):
async with self.app.install_lock_file.shared(): async with self.app.install_lock_file.shared():
self._crash_reports = {} self._crash_reports = {}
if isinstance(self.ui.body, ProbingFailed): if isinstance(self.ui.body, ProbingFailed):
self.ui.set_body(SlowProbing(self)) self.ui.set_body(SlowProbing(self))
schedule_task(self._wait_for_probing()) schedule_task(self._wait_for_probing())
for (restricted, kind) in [ for (restricted, kind) in [
(False, ErrorReportKind.BLOCK_PROBE_FAIL), (False, ErrorReportKind.BLOCK_PROBE_FAIL),
(True, ErrorReportKind.DISK_PROBE_FAIL), (True, ErrorReportKind.DISK_PROBE_FAIL),
]: ]:
try: try:
desc = "restricted={}".format(restricted) await self._probe_once_task.start(
with context.child("probe_once", desc): context=context, restricted=restricted)
await self._probe_once_task.start(restricted) # We wait on the task directly here, not
# We wait on the task directly here, not # self._probe_once_task.wait as if _probe_once_task
# self._probe_once_task.wait as if _probe_once_task # gets cancelled, we should be cancelled too.
# gets cancelled, we should be cancelled too. await asyncio.wait_for(self._probe_once_task.task, 15.0)
await asyncio.wait_for( except asyncio.CancelledError:
self._probe_once_task.task, 15.0) # asyncio.CancelledError is a subclass of Exception in
except asyncio.CancelledError: # Python 3.6 (sadface)
# asyncio.CancelledError is a subclass of Exception in raise
# Python 3.6 (sadface) except Exception:
raise block_discover_log.exception(
except Exception: "block probing failed restricted=%s", restricted)
block_discover_log.exception( report = self.app.make_apport_report(
"block probing failed restricted=%s", restricted) kind, "block probing", interrupt=False)
report = self.app.make_apport_report( self._crash_reports[restricted] = report
kind, "block probing", interrupt=False) continue
self._crash_reports[restricted] = report break
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) log.debug("self.ai_data = %s", self.ai_data)
if 'layout' in self.ai_data: if 'layout' in self.ai_data:
layout = self.ai_data['layout'] layout = self.ai_data['layout']
with self.context.child("applying_autoinstall"): meth = getattr(self, "guided_" + layout['name'])
meth = getattr(self, "guided_" + layout['name']) disk = self.model.disk_for_match(
disk = self.model.disk_for_match( self.model.all_disks(),
self.model.all_disks(), layout.get("match", {'size': 'largest'}))
layout.get("match", {'size': 'largest'})) meth(disk)
meth(disk)
elif 'config' in self.ai_data: 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.grub = self.ai_data.get('grub', {})
self.model.swap = self.ai_data.get('swap') self.model.swap = self.ai_data.get('swap')

View File

@ -45,7 +45,7 @@ class IdentityController(SubiquityController):
self.model.add_user(data) self.model.add_user(data)
@with_context() @with_context()
async def apply_autoinstall_config(self, context): async def apply_autoinstall_config(self, context=None):
if not self.model.user: if not self.model.user:
if 'user-data' not in self.app.autoinstall_config: if 'user-data' not in self.app.autoinstall_config:
raise Exception("no identity data provided") raise Exception("no identity data provided")

View File

@ -241,7 +241,7 @@ class InstallProgressController(SubiquityController):
return curtin_cmd return curtin_cmd
@with_context(description="umounting /target dir") @with_context(description="umounting /target dir")
async def unmount_target(self, context, target): async def unmount_target(self, *, context, target):
cmd = [ cmd = [
sys.executable, '-m', 'curtin', 'unmount', sys.executable, '-m', 'curtin', 'unmount',
'-t', target, '-t', target,
@ -254,7 +254,7 @@ class InstallProgressController(SubiquityController):
@with_context( @with_context(
description="installing system", level="INFO", childlevel="DEBUG") description="installing system", level="INFO", childlevel="DEBUG")
async def curtin_install(self, context): async def curtin_install(self, *, context):
log.debug('curtin_install') log.debug('curtin_install')
self.install_state = InstallState.RUNNING self.install_state = InstallState.RUNNING
self.curtin_event_contexts[''] = context self.curtin_event_contexts[''] = context
@ -288,7 +288,7 @@ class InstallProgressController(SubiquityController):
pass pass
@with_context() @with_context()
async def install(self, context): async def install(self, *, context):
context.set('is-install-context', True) context.set('is-install-context', True)
try: try:
await asyncio.wait( await asyncio.wait(
@ -297,16 +297,17 @@ class InstallProgressController(SubiquityController):
await self.confirmation.wait() await self.confirmation.wait()
if os.path.exists(self.model.target): 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( await asyncio.wait(
{e.wait() for e in self.model.postinstall_events}) {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.ui.set_header(_("Installation complete!"))
self.progress_view.set_status(_("Finished install!")) self.progress_view.set_status(_("Finished install!"))
@ -314,7 +315,7 @@ class InstallProgressController(SubiquityController):
if self.model.network.has_network: if self.model.network.has_network:
self.progress_view.update_running() self.progress_view.update_running()
await self.run_unattended_upgrades(context) await self.run_unattended_upgrades(context=context)
self.progress_view.update_done() self.progress_view.update_done()
except Exception: except Exception:
@ -326,7 +327,7 @@ class InstallProgressController(SubiquityController):
await self.install_task await self.install_task
self.app.next_screen() self.app.next_screen()
async def drain_curtin_events(self, context): async def drain_curtin_events(self, *, context):
waited = 0.0 waited = 0.0
while self.progress_view.ongoing and waited < 5.0: while self.progress_view.ongoing and waited < 5.0:
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
@ -337,30 +338,29 @@ class InstallProgressController(SubiquityController):
@with_context( @with_context(
description="final system configuration", level="INFO", description="final system configuration", level="INFO",
childlevel="DEBUG") childlevel="DEBUG")
async def postinstall(self, context): async def postinstall(self, *, context):
autoinstall_path = os.path.join( autoinstall_path = os.path.join(
self.app.root, 'var/log/installer/autoinstall-user-data') self.app.root, 'var/log/installer/autoinstall-user-data')
autoinstall_config = "#cloud-config\n" + yaml.dump( autoinstall_config = "#cloud-config\n" + yaml.dump(
{"autoinstall": self.app.make_autoinstall()}) {"autoinstall": self.app.make_autoinstall()})
write_file(autoinstall_path, autoinstall_config, mode=0o600) write_file(autoinstall_path, autoinstall_config, mode=0o600)
await self.configure_cloud_init(context) await self.configure_cloud_init(context=context)
packages = [] packages = []
if self.model.ssh.install_server: if self.model.ssh.install_server:
packages = ['openssh-server'] packages = ['openssh-server']
packages.extend(self.app.base_model.packages) packages.extend(self.app.base_model.packages)
for package in packages: for package in packages:
subcontext = context.child( await self.install_package(context=context, package=package)
"install_{}".format(package), await self.restore_apt_config(context=context)
"installing {}".format(package))
with subcontext:
await self.install_package(package)
await self.restore_apt_config(context)
@with_context(description="configuring cloud-init") @with_context(description="configuring cloud-init")
async def configure_cloud_init(self, context): async def configure_cloud_init(self, context):
await run_in_thread(self.model.configure_cloud_init) 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: if self.opts.dry_run:
cmd = ["sleep", str(2/self.app.scale_factor)] cmd = ["sleep", str(2/self.app.scale_factor)]
else: else:

View File

@ -155,9 +155,9 @@ class NetworkController(NetworkController, SubiquityController):
self.model.has_network = bool( self.model.has_network = bool(
self.network_event_receiver.default_routes) self.network_event_receiver.default_routes)
async def _apply_config(self, context=None, *, silent): async def _apply_config(self, *, context=None, silent=False):
try: try:
await super()._apply_config(context, silent=silent) await super()._apply_config(context=context, silent=silent)
except asyncio.CancelledError: except asyncio.CancelledError:
# asyncio.CancelledError is a subclass of Exception in # asyncio.CancelledError is a subclass of Exception in
# Python 3.6 (sadface) # Python 3.6 (sadface)

View File

@ -43,7 +43,7 @@ class ProxyController(SubiquityController):
self.signal.emit_signal('network-proxy-set') self.signal.emit_signal('network-proxy-set')
@with_context() @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 # XXX want to wait until signal sent by .start() has been seen
# by everything; don't have a way to do that today. # by everything; don't have a way to do that today.
pass pass

View File

@ -36,25 +36,25 @@ class RebootController(SubiquityController):
def interactive(self): def interactive(self):
return self.app.interactive() return self.app.interactive()
async def copy_logs_to_target(self): @with_context()
with self.context.child("copy_logs_to_target"): async def copy_logs_to_target(self, context):
if self.opts.dry_run and 'copy-logs-fail' in self.app.debug_flags: if self.opts.dry_run and 'copy-logs-fail' in self.app.debug_flags:
raise PermissionError() raise PermissionError()
target_logs = os.path.join( target_logs = os.path.join(
self.app.base_model.target, 'var/log/installer') self.app.base_model.target, 'var/log/installer')
if self.opts.dry_run: if self.opts.dry_run:
os.makedirs(target_logs, exist_ok=True) os.makedirs(target_logs, exist_ok=True)
else: 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( await arun_command(
['cp', '-aT', '/var/log/installer', target_logs]) ['journalctl', '-b'],
journal_txt = os.path.join(target_logs, 'installer-journal.txt') stdout=output, stderr=subprocess.STDOUT)
try: except Exception:
with open(journal_txt, 'w') as output: log.exception("saving journal failed")
await arun_command(
['journalctl', '-b'],
stdout=output, stderr=subprocess.STDOUT)
except Exception:
log.exception("saving journal failed")
def reboot(self): def reboot(self):
if self.opts.dry_run: if self.opts.dry_run:
@ -66,7 +66,7 @@ class RebootController(SubiquityController):
@with_context() @with_context()
async def apply_autoinstall_config(self, context): async def apply_autoinstall_config(self, context):
await self.copy_logs_to_target() await self.copy_logs_to_target(context=context)
self.reboot() self.reboot()
async def _run(self): async def _run(self):

View File

@ -97,7 +97,7 @@ class RefreshController(SubiquityController):
return return
if self.check_state != CheckState.AVAILABLE: if self.check_state != CheckState.AVAILABLE:
return return
change_id = await self.start_update(context) change_id = await self.start_update(context=context)
while True: while True:
try: try:
change = await self.get_progress(change_id) change = await self.get_progress(change_id)

View File

@ -20,6 +20,7 @@ import requests.exceptions
from subiquitycore.async_helpers import ( from subiquitycore.async_helpers import (
schedule_task, schedule_task,
) )
from subiquitycore.context import with_context
from subiquitycore.controller import ( from subiquitycore.controller import (
Skip, Skip,
) )
@ -62,34 +63,34 @@ class SnapdSnapInfoLoader:
while self.pending_snaps: while self.pending_snaps:
snap = self.pending_snaps.pop(0) snap = self.pending_snaps.pop(0)
task = self.tasks[snap] = schedule_task( task = self.tasks[snap] = schedule_task(
self._fetch_info_for_snap(snap)) self._fetch_info_for_snap(snap=snap))
await task await task
async def _load_list(self): @with_context(name="list")
with self.context.child("list"): async def _load_list(self, context=None):
try: try:
result = await self.snapd.get( result = await self.snapd.get(
'v2/find', section=self.store_section) 'v2/find', section=self.store_section)
except requests.exceptions.RequestException: except requests.exceptions.RequestException:
log.exception("loading list of snaps failed") log.exception("loading list of snaps failed")
self.failed = True self.failed = True
return return
self.model.load_find_data(result) self.model.load_find_data(result)
self.snap_list_fetched = True self.snap_list_fetched = True
def stop(self): def stop(self):
if self.main_task is not None: if self.main_task is not None:
self.main_task.cancel() self.main_task.cancel()
async def _fetch_info_for_snap(self, snap): @with_context(name="fetch/{snap.name}")
with self.context.child("fetch").child(snap.name): async def _fetch_info_for_snap(self, snap, context=None):
try: try:
data = await self.snapd.get('v2/find', name=snap.name) data = await self.snapd.get('v2/find', name=snap.name)
except requests.exceptions.RequestException: except requests.exceptions.RequestException:
log.exception("loading snap info failed") log.exception("loading snap info failed")
# XXX something better here? # XXX something better here?
return return
self.model.load_info_data(data) self.model.load_info_data(data)
def get_snap_list_task(self): def get_snap_list_task(self):
return self.tasks[None] return self.tasks[None]
@ -98,7 +99,8 @@ class SnapdSnapInfoLoader:
if snap not in self.tasks: if snap not in self.tasks:
if snap in self.pending_snaps: if snap in self.pending_snaps:
self.pending_snaps.remove(snap) 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] return self.tasks[snap]

View File

@ -17,6 +17,7 @@ import logging
import subprocess import subprocess
from subiquitycore.async_helpers import schedule_task from subiquitycore.async_helpers import schedule_task
from subiquitycore.context import with_context
from subiquitycore import utils from subiquitycore import utils
from subiquity.controller import SubiquityController from subiquity.controller import SubiquityController
@ -94,7 +95,10 @@ class SSHController(SubiquityController):
raise subprocess.CalledProcessError(cp.returncode, cmd) raise subprocess.CalledProcessError(cp.returncode, cmd)
return cp 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) ssh_import_id = "{ssh_import_id}:{import_username}".format(**user_spec)
with self.context.child("ssh_import_id", ssh_import_id): with self.context.child("ssh_import_id", ssh_import_id):
try: try:
@ -127,7 +131,8 @@ class SSHController(SubiquityController):
user_spec, ssh_import_id, key_material, fingerprints) user_spec, ssh_import_id, key_material, fingerprints)
def fetch_ssh_keys(self, user_spec): 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): def done(self, result):
log.debug("SSHController.done next_screen result=%s", result) log.debug("SSHController.done next_screen result=%s", result)

View File

@ -15,6 +15,8 @@
import asyncio import asyncio
import enum import enum
import functools
import inspect
class Status(enum.Enum): class Status(enum.Enum):
@ -117,12 +119,30 @@ def with_context(name=None, description="", **context_kw):
if name is None: if name is None:
name = meth.__name__ name = meth.__name__
async def decorated(self, context=None, *args, **kw): def convargs(self, kw):
context = kw.get('context')
if context is None: if context is None:
context = self.context context = self.context
manager = context.child( kw['context'] = context.child(
name, description=description.format(**kw), **context_kw) name=name.format(**kw),
with manager as subcontext: description=description.format(**kw),
await meth(self, subcontext, *args, **kw) **context_kw)
return decorated 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 return decorate

View File

@ -317,7 +317,7 @@ class NetworkController(BaseController):
return os.path.join(self.root, 'etc/netplan', netplan_config_file_name) return os.path.join(self.root, 'etc/netplan', netplan_config_file_name)
def apply_config(self, context=None, silent=False): 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): async def _down_devs(self, devs):
for dev in devs: for dev in devs:
@ -358,7 +358,7 @@ class NetworkController(BaseController):
@with_context( @with_context(
name="apply_config", description="silent={silent}", level="INFO") 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_delete = []
devs_to_down = [] devs_to_down = []
dhcp_device_versions = [] dhcp_device_versions = []