Implement disk quota checking.

Signed-off-by: Edward Z. Yang <ezyang@mit.edu>

diff --git a/TODO b/TODO
index dc3c3227628438486345e143d76031acfcd765a7..21698b2e58e892db278ed26133d7b6f3e5b471be 100644 (file)
--- a/TODO
+++ b/TODO
@@ -7,10 +7,6 @@ TODO NOW:
       does not seem to have been a problem in practice)
     - Custom merge algo: check if it's got extra \r's in the file,
       and dos2unix it if it does, before performing the merge
-    - `vos exa` in order to check what a person's quota is.  We can
-      figure out roughly how big the upgrade is going to be by
-      doing a size comparison of the tars: `git pull` MUST NOT
-      fail, otherwise things are left conflicted, and not easy to fix.
     - Prune -7 call errors and automatically reprocess them (with a
       strike out counter of 3)--this requires better error parsing
 

diff --git a/wizard/command/upgrade.py b/wizard/command/upgrade.py
index 18e7f9efe837d3d7990015fc7c2e8714bd858cc9..a8767dae9370a6d8303f2aa5a7ee2ea5c1acd9fd 100644 (file)
--- a/wizard/command/upgrade.py
+++ b/wizard/command/upgrade.py
@@ -8,7 +8,9 @@ import errno
 import tempfile
 import itertools
 
-from wizard import app, command, deploy, shell, util
+from wizard import app, command, deploy, scripts, shell, util
+
+kib_buffer = 1024 * 1024 * 10 # 10 MiB we will always leave available
 
 def main(argv, baton):
     options, args = parse_args(argv, baton)
@@ -49,6 +51,13 @@ def main(argv, baton):
         d.verifyVersion()
         if not options.dry_run:
             d.verifyWeb()
+        # perform database backup
+        if not options.dry_run:
+            backup = d.backup(options)
+        # we do this before checking quota
+        kib_usage, kib_limit = scripts.get_quota_usage_and_limit()
+        if kib_limit and (kib_limit - kib_usage) < kib_buffer: # 10 mebibytes
+            raise QuotaTooLow
         repo = d.application.repository(options.srv_path)
         version = calculate_newest_version(sh, repo)
         if version == d.app_version.scripts_tag and not options.force:
@@ -80,7 +89,7 @@ def main(argv, baton):
             open(".git/WIZARD_UPGRADE_VERSION", "w").write(version)
             open(".git/WIZARD_PARENTS", "w").write("%s\n%s" % (user_commit, next_commit))
             if options.log_file: open(".git/WIZARD_LOG_FILE", "w").write(options.log_file)
-            perform_merge(sh, repo, d, version, use_shm)
+            perform_merge(sh, repo, d, version, use_shm, kib_limit and kib_limit - kib_usage or None)
     # variables: version, user_commit, next_commit, temp_wc_dir
     with util.ChangeDirectory(temp_wc_dir):
         try:
@@ -106,8 +115,6 @@ def main(argv, baton):
     if options.dry_run:
         logging.info("Dry run, bailing.  See results at %s" % temp_wc_dir)
         return
-    # perform database backup
-    backup = d.backup(options)
     # XXX: frob .htaccess to make site inaccessible
     with util.IgnoreKeyboardInterrupts():
         with util.LockDirectory(".scripts-upgrade-lock"):
@@ -189,7 +196,8 @@ def perform_tmp_clone(sh, use_shm):
     sh.call("git", "clone", "-q", "--shared", ".", temp_wc_dir)
     return temp_dir, temp_wc_dir
 
-def perform_merge(sh, repo, d, version, use_shm):
+def perform_merge(sh, repo, d, version, use_shm, kib_avail):
+    # Note: avail_quota == None means unlimited
     # naive merge algorithm:
     # sh.call("git", "merge", "-m", message, "scripts/master")
     # crazy merge algorithm:
@@ -214,6 +222,7 @@ def perform_merge(sh, repo, d, version, use_shm):
     user_virtual_commit = sh.eval("git", "commit-tree", user_tree,
             "-p", base_virtual_commit, input="", log=True)
     sh.call("git", "checkout", user_virtual_commit, "--")
+    pre_usage = scripts.get_disk_usage()
     try:
         sh.call("git", "merge", next_virtual_commit)
     except shell.CallError as e:
@@ -239,8 +248,14 @@ def perform_merge(sh, repo, d, version, use_shm):
             shutil.move(curdir, newdir)
             shutil.rmtree(os.path.dirname(curdir))
             curdir = os.path.join(newdir, "repo")
+            os.chdir(curdir)
         print "%d %s" % (conflicts, curdir)
         raise MergeFailed
+    finally:
+        post_usage = scripts.get_disk_usage()
+        kib_delta = (post_usage - pre_usage) / 1024
+        if kib_avail is not None and kib_delta - kib_avail < kib_buffer:
+            raise QuotaTooLow
 
 def parse_args(argv, baton):
     usage = """usage: %prog upgrade [ARGS] [DIR]
@@ -268,6 +283,14 @@ class Error(command.Error):
     """Base exception for all exceptions raised by upgrade"""
     pass
 
+class QuotaTooLow(Error):
+    def __str__(self):
+        return """
+
+ERROR: The locker quota was too low to complete the autoinstall
+upgrade.
+"""
+
 class AlreadyUpgraded(Error):
     def __str__(self):
         return """

diff --git a/wizard/scripts.py b/wizard/scripts.py
index 8122c78ec5e3c9e5ce728de741b5eead0ce5dd73..654fe6865ae1e4a16d3f34c07dfd7dfd794b768c 100644 (file)
--- a/wizard/scripts.py
+++ b/wizard/scripts.py
@@ -43,3 +43,45 @@ def get_web_host_and_path(dir=None):
         return None
     return (util.get_dir_owner(dir) + ".scripts.mit.edu", web_path)
 
+def get_quota_usage_and_limit(dir=None):
+    """
+    Returns a tuple (quota usage, quota limit).  Works only for scripts
+    servers.  Values are in KiB.  Returns ``(0, None)`` if we couldn't figure it out.
+    """
+    # XXX: The correct way is to implement Python modules implementing
+    # bindings for all the appropriate interfaces
+    def parse_last_quote(ret):
+        return ret.rstrip('\'').rpartition('\'')[2]
+    if dir is None:
+        dir = os.getcwd()
+    sh = shell.Shell()
+    try:
+        cell = parse_last_quote(sh.eval("fs", "whichcell", "-path", dir))
+    except shell.CallError:
+        return (0, None)
+    mount = None
+    while dir:
+        try:
+            volume = parse_last_quote(sh.eval("fs", "lsmount", dir))[1:]
+            break
+        except shell.CallError:
+            dir = os.path.dirname(dir)
+    print volume
+    if not volume: return (0, None)
+    try:
+        result = sh.eval("vos", "examine", "-id", volume, "-cell", cell).splitlines()
+    except shell.CallError:
+        return (0, None)
+    usage = int(result[0].split()[3])
+    limit = int(result[3].split()[1]) # XXX: FRAGILE
+    return (usage, limit)
+
+# XXX: Possibly in the wrong module
+def get_disk_usage(dir=None):
+    """
+    Recursively determines the disk usage of a directory, excluding
+    .git directories.  Value is in bytes.
+    """
+    if dir is None: dir = os.getcwd()
+    return int(shell.Shell().eval("du", "--exclude=.git", "-s", dir).split()[0])
+