Make MINION report Lorry's disk usage after a job finishes

1 files changed, 51 insertions, 0 deletions

diff --git a/lorry-controller-minion b/lorry-controller-minion
index 7b0013a..0c86a03 100755
--- a/lorry-controller-minion
+++ b/lorry-controller-minion
@@ -70,6 +70,12 @@ class MINION(cliapp.Application):
             metavar='CMD',
             default='lorry')
 
+        self.settings.string(
+            ['lorry-working-area'],
+            'where will Lorry put its files?',
+            metavar='DIR',
+            default='/home/lorry/working-area')
+
     def process_args(self, args):
         logging.info('Starting MINION')
 
@@ -197,11 +203,17 @@ class MINION(cliapp.Application):
         logging.debug(
             'Updating WEBAPP about running job %s', job_spec['job_id'])
 
+        if exit is None:
+            disk_usage = None
+        else:
+            disk_usage = self.get_lorry_disk_usage(job_spec)
+
         params = urllib.urlencode({
                 'job_id': job_spec['job_id'],
                 'exit': 'no' if exit is None else exit,
                 'stdout': stdout,
                 'stderr': stderr,
+                'disk_usage': disk_usage,
                 })
 
         try:
@@ -238,4 +250,43 @@ class MINION(cliapp.Application):
 
         return response_body
 
+    def get_lorry_disk_usage(self, job_spec):
+        dirname = os.path.join(
+            self.settings['lorry-working-area'],
+            self.escape_lorry_area_basename(job_spec['path']))
+        return self.disk_usage_by_dir(dirname)
+
+    def escape_lorry_area_basename(self, basename):
+        # FIXME: This code should be kept in sync with the respective
+        # code in lorry, or, better, we would import the code from
+        # Lorry directly.
+
+        assert '\0' not in basename
+        # We escape slashes as underscores.
+        return '_'.join(basename.split('/'))
+
+    def disk_usage_by_dir(self, dirname):
+        exit, out, err = cliapp.runcmd_unchecked(['du', '-sk', dirname])
+        if exit:
+            logging.error('du -sk %s failed: %r', dirname, err)
+            return 0
+
+        lines = out.splitlines()
+        if not lines:
+            logging.warning('no output from du')
+            return 0
+
+        words = lines[-1].split()
+        if not words:
+            logging.warning('last line of du output is empty')
+            return 0
+
+        kibibyte = 1024
+        try:
+            return int(words[0]) * kibibyte
+        except ValueError:
+            logging.warning('error converting %r to string' % words[0])
+            return 0
+        
+
 MINION().run()