Refactor: use a 'with Timer' construct for timers

1 files changed, 90 insertions, 91 deletions

diff --git a/contractor b/contractor
index 58625af..38c4b39 100755
--- a/contractor
+++ b/contractor
@@ -57,98 +57,92 @@ class ContractorApplication(cliapp.Application):
         self.output.write('{}\n'.format(json.dumps(bs.as_dict(), indent=4)))
 
     def cmd_build(self, args):
-        timer = Timer(self.verbose)
-        overall = Timer(self.verbose)
-        dest = self.manager_destination()
-        manager = RemoteServer(dest, verbose=self.verbose)
-
         self.verbose('building using spec at {}'.format(args[0]))
         bs = self.load_build_spec(args[0])
+        dest = self.manager_destination()
+        manager = RemoteServer(dest, verbose=self.verbose)
 
-        self.upload_worker_image(bs.worker_image(), dest)
-        timer.report('upload-worker-image')
-
-        # Do the minimum needed to start worker VM. The VM takes a
-        # while to boot and we can do other things while that happens.
-        execs = [
-            GetCPUCount(),
-        ]
-        er = self.exec_quietly(manager, *execs)
-        timer.report('setup')
-
-        # Find number of CPUs.
-        cpus = 1
-        for line in er.stdout.decode('UTF8').splitlines():
-            if line.startswith('CPU(s):'):
-                cpus = int(line.split()[-1])
-
-        execs = [
-            DestroyWorkerVM(),
-            UndefineWorkerVM(),
-            CopyWorkerImage(),
-            StartGuestNetworking(),
-            CreateWorkerVM(cpus),
-            TryUnmountWS(),
-            MountWS(),
-            ChownWS(),
-        ]
-        self.exec_quietly(manager, *execs)
-        timer.report('setup')
-
-        self.verbose('setting up workspace on worker')
-
-        ws = bs.workspace()
-        if os.path.exists(ws):
-            self.sync_to_workspace(ws, dest, '.')
-            timer.report('upload-saved-workspace')
-
-        execs = [
-            Mkdir('/mnt/src', owner=WORKER_UID, group=WORKER_GID),
-        ]
-        self.exec_quietly(manager, *execs)
-        src = bs.source()
-        self.sync_to_workspace(src, dest, 'src')
-        timer.report('upload-source')
-
-        execs = [
-            UnmountWS(),
-            WorkerIP(),
-        ]
-        er = self.exec_quietly(manager, *execs)
-        worker_ip = er.stdout.decode('UTF8').strip()
-        timer.report('wait-for-worker')
-
-        self.exec_quietly(manager, AttachWS())
-        self.verbose('attached')
-
-        worker = OnWorker(
-            dest, 'worker@{}'.format(worker_ip), verbose=self.verbose)
-        self.exec_quietly(worker, Mkdir('/workspace'), MountWSonWorker())
-
-        ansible = bs.ansible()
-        if ansible:
-            self.exec_verbosely(manager, Ansible(ansible, worker_ip))
-
-        execs = [
-            Chdir('/workspace/src'),
-            Build(bs.build()),
-        ]
-        self.exec_verbosely(worker, *execs)
-        timer.report('build')
-
-        execs = [
-            ShutdownWorkerVM(),
-            MountWS(),
-        ]
-        self.exec_quietly(manager, *execs)
-        timer.report('shutdown-worker')
-
-        if ws:
-            self.verbose('saving workspace to {}'.format(ws))
-            self.sync_from_workspace(dest, ws)
-            timer.report('save-workspace')
+        with Timer(self.verbose, 'complete-run'):
+            with Timer(self.verbose, 'upload-worker-image'):
+                self.upload_worker_image(bs.worker_image(), dest)
+
+            # Do the minimum needed to start worker VM. The VM takes a
+            # while to boot and we can do other things while that
+            # happens.
+            with Timer(self.verbose, 'start-worker'):
+                # Find number of CPUs.
+                er = self.exec_quietly(manager, GetCPUCount())
+                cpus = 1
+                for line in er.stdout.decode('UTF8').splitlines():
+                    if line.startswith('CPU(s):'):
+                        cpus = int(line.split()[-1])
+
+                execs = [
+                    DestroyWorkerVM(),
+                    UndefineWorkerVM(),
+                    CopyWorkerImage(),
+                    StartGuestNetworking(),
+                    CreateWorkerVM(cpus),
+                    TryUnmountWS(),
+                    MountWS(),
+                    ChownWS(),
+                ]
+                self.exec_quietly(manager, *execs)
+
+            with Timer(self.verbose, 'upload-saved-workspace'):
+                ws = bs.workspace()
+                if os.path.exists(ws):
+                    self.sync_to_workspace(ws, dest, '.')
+
+            with Timer(self.verbose, 'upload-source'):
+                self.exec_quietly(
+                    manager, Mkdir(
+                        '/mnt/src', owner=WORKER_UID, group=WORKER_GID))
+                src = bs.source()
+                self.sync_to_workspace(src, dest, 'src')
+
+            with Timer(self.verbose, 'wait-for-worker-to-be-available'):
+                execs = [
+                    UnmountWS(),
+                    WorkerIP(),
+                ]
+                er = self.exec_quietly(manager, *execs)
+                worker_ip = er.stdout.decode('UTF8').strip()
+                timer.report('wait-for-worker')
+
+            with Timer(self.verbose, 'prepare-workspace-worker'):
+                self.exec_quietly(manager, AttachWS())
+                self.verbose('attached')
+
+                worker = OnWorker(
+                    dest, 'worker@{}'.format(worker_ip), verbose=self.verbose)
+                self.exec_quietly(worker, Mkdir('/workspace'), MountWSonWorker())
+
+            with Timer(self.verbose, 'prepare-worker-with-ansible'):
+                ansible = bs.ansible()
+                if ansible:
+                    self.exec_verbosely(manager, Ansible(ansible, worker_ip))
+
+            with Timer(self.verbose, 'build'):
+                execs = [
+                    Chdir('/workspace/src'),
+                    Build(bs.build()),
+                ]
+                self.exec_verbosely(worker, *execs)
+
+            with Timer(self.verbose, 'shutdown-worker'):
+                execs = [
+                    ShutdownWorkerVM(),
+                    MountWS(),
+                ]
+                self.exec_quietly(manager, *execs)
+
+            with Timer(self.verbose, 'save-workspace'):
+                if ws:
+                    self.verbose('saving workspace to {}'.format(ws))
+                    self.sync_from_workspace(dest, ws)
+                    timer.report('save-workspace')
 
-        overall.report('complete-run')
         self.verbose('build finished OK')
 
     def load_build_spec(self, filename):
@@ -617,15 +611,20 @@ class SpecMissingKey(Exception):
 
 class Timer:
 
-    def __init__(self, report):
+    def __init__(self, report, title):
         self._report = report
+        self._title = title
+        self._prev = None
+
+    def __enter__(self):
         self._prev = time.time()
 
-    def report(self, msg):
+    def __exit__(self, exctype, exc, tb):
         now = time.time()
         duration = now - self._prev
         self._prev = now
-        self._report('time: {:.1f} s {}'.format(duration, msg))
+        self._report('time: {:.1f} s {}'.format(duration, self._title))
+        return exctype
 
 
 ContractorApplication().run()