Reworked stage paths to allow %u for username. Added stage test.

8 files changed, 385 insertions, 66 deletions

diff --git a/lib/spack/spack/cmd/checksum.py b/lib/spack/spack/cmd/checksum.py
index 7dd214e6aa..f50e231104 100644
--- a/lib/spack/spack/cmd/checksum.py
+++ b/lib/spack/spack/cmd/checksum.py
@@ -36,8 +36,7 @@ def checksum(parser, args):
     if not versions:
         versions = pkg.fetch_available_versions()[:args.number]
         if not versions:
-            tty.die("Could not fetch any available versions for %s."
-                    % pkg.name)
+            tty.die("Could not fetch any available versions for %s." % pkg.name)
 
     versions.sort()
     versions.reverse()
@@ -48,7 +47,7 @@ def checksum(parser, args):
 
     hashes = []
     for url, version in zip(urls, versions):
-        stage = Stage("checksum-%s-%s" % (pkg.name, version), url)
+        stage = Stage(url)
         try:
             stage.fetch()
             hashes.append(md5(stage.archive_file))
diff --git a/lib/spack/spack/cmd/create.py b/lib/spack/spack/cmd/create.py
index 79cd9c6b17..cc3274f70e 100644
--- a/lib/spack/spack/cmd/create.py
+++ b/lib/spack/spack/cmd/create.py
@@ -57,7 +57,7 @@ def create(parser, args):
 
     # make a stage and fetch the archive.
     try:
-        stage = Stage("spack-create/%s-%s" % (name, version), url)
+        stage = Stage(url)
         archive_file = stage.fetch()
     except spack.FailedDownloadException, e:
         tty.die(e.message)
diff --git a/lib/spack/spack/globals.py b/lib/spack/spack/globals.py
index 7b746a0734..a2a14cfad1 100644
--- a/lib/spack/spack/globals.py
+++ b/lib/spack/spack/globals.py
@@ -54,10 +54,12 @@ curl = which("curl", required=True)
 use_tmp_stage = True
 
 # Locations to use for staging and building, in order of preference
-# Spack will try to create stage directories in <tmp_dir>/<username>
-# if one of these tmp_dirs exists.  Otherwise it'll use a default
-# location per the python implementation of tempfile.mkdtemp().
-tmp_dirs = ['/nfs/tmp2', '/var/tmp', '/tmp']
+# Use a %u to add a username to the stage paths here, in case this
+# is a shared filesystem.  Spack will use the first of these paths
+# that it can create.
+tmp_dirs = ['/nfs/tmp2/%u/spack-stage',
+            '/var/tmp/%u/spcak-stage',
+            '/tmp/%u/spack-stage']
 
 #
 # SYS_TYPE to use for the spack installation.
diff --git a/lib/spack/spack/package.py b/lib/spack/spack/package.py
index 8cbb3f0eaf..049d193eef 100644
--- a/lib/spack/spack/package.py
+++ b/lib/spack/spack/package.py
@@ -313,7 +313,8 @@ class Package(object):
             self.versions = VersionList(self.versions)
 
         # stage used to build this package.
-        self.stage = Stage("%s-%s" % (self.name, self.version), self.url)
+        # TODO: hash the concrete spec and use that as the stage name.
+        self.stage = Stage(self.url, "%s-%s" % (self.name, self.version))
 
         # Set a default list URL (place to find available versions)
         if not hasattr(self, 'list_url'):
diff --git a/lib/spack/spack/stage.py b/lib/spack/spack/stage.py
index 9bf9584f57..d38413d155 100644
--- a/lib/spack/spack/stage.py
+++ b/lib/spack/spack/stage.py
@@ -2,18 +2,15 @@ import os
 import re
 import shutil
 import tempfile
-import getpass
 
 import spack
 import spack.error as serr
-import tty
+import spack.tty as tty
 
-class FailedDownloadError(serr.SpackError):
-    """Raised wen a download fails."""
-    def __init__(self, url):
-        super(FailedDownloadError, self).__init__(
-            "Failed to fetch file from URL: " + url)
-        self.url = url
+from spack.util.filesystem import *
+from spack.util.compression import decompressor_for
+
+STAGE_PREFIX = 'spack-stage-'
 
 
 class Stage(object):
@@ -31,16 +28,76 @@ class Stage(object):
        If spack.use_tmp_stage is True, spack will attempt to create stages
        in a tmp directory.  Otherwise, stages are created directly in
        spack.stage_path.
+
+       There are two kinds of stages: named and unnamed.  Named stages can
+       persist between runs of spack, e.g. if you fetched a tarball but
+       didn't finish building it, you won't have to fetch it again.
+
+       Unnamed stages are created using standard mkdtemp mechanisms or
+       similar, and are intended to persist for only one run of spack.
     """
-    def __init__(self, path, url):
+
+    def __init__(self, url, name=None):
         """Create a stage object.
            Parameters:
-             path    Relative path from the stage root to where the stage will
-                     be created.
              url     URL of the archive to be downloaded into this stage.
+
+             name    If a name is provided, then this stage is a named stage
+                     and will persist between runs (or if you construct another
+                     stage object later).  If name is not provided, then this
+                     stage will be given a unique name automatically.
         """
-        self.path = os.path.join(spack.stage_path, path)
+        self.tmp_root = find_tmp_root()
         self.url = url
+        self.name = name
+        self.path = None   # This will be set after setup is called.
+
+
+    def _cleanup_dead_links(self):
+        """Remove any dead links in the stage directory."""
+        for file in os.listdir(spack.stage_path):
+            path = new_path(spack.stage_path, file)
+            if os.path.islink(path):
+                real_path = os.path.realpath(path)
+                if not os.path.exists(path):
+                    os.unlink(path)
+
+
+    def _need_to_create_path(self):
+        """Makes sure nothing weird has happened since the last time we
+           looked at path.  Returns True if path already exists and is ok.
+           Returns False if path needs to be created.
+        """
+        # Path doesn't exist yet.  Will need to create it.
+        if not os.path.exists(self.path):
+            return True
+
+        # Path exists but points at something else.  Blow it away.
+        if not os.path.isdir(self.path):
+            os.unlink(self.path)
+            return True
+
+        # Path looks ok, but need to check the target of the link.
+        if os.path.islink(self.path):
+            real_path = os.path.realpath(self.path)
+
+            if spack.use_tmp_stage:
+                # If we're using a tmp dir, it's a link, and it points at the right spot,
+                # then keep it.
+                if (os.path.commonprefix((real_path, self.tmp_root)) == self.tmp_root
+                    and os.path.exists(real_path)):
+                    return False
+                else:
+                    # otherwise, just unlink it and start over.
+                    os.unlink(self.path)
+                    return True
+
+            else:
+                # If we're not tmp mode, then it's a link and we want a directory.
+                os.unlink(self.path)
+                return True
+
+        return False
 
 
     def setup(self):
@@ -54,53 +111,39 @@ class Stage(object):
            create a stage.  If there is no valid location in tmp_dirs, fall
            back to making the stage inside spack.stage_path.
         """
-        # If we're using a stage in tmp that has since been deleted,
-        # remove the stale symbolic link.
-        if os.path.islink(self.path):
-            real_path = os.path.realpath(self.path)
-            if not os.path.exists(real_path):
-                os.unlink(self.path)
-
-        # If the user switched stage modes, destroy the old stage and
-        # start over.  We could move the old archive, but that seems
-        # like a pain when we could just fetch it again.
-        if spack.use_tmp_stage:
-            if not os.path.islink(self.path):
-                self.destroy()
-        else:
-            if os.path.islink(self.path):
-                self.destroy()
-
-        # Make sure that the stage is actually a directory.  Something
-        # is seriously wrong if it's not.
-        if os.path.exists(self.path):
-            if not os.path.isdir(self.path):
-                tty.die("Stage path %s is not a directory!" % self.path)
-        else:
-            # Create the top-level stage directory
-            spack.mkdirp(spack.stage_path)
-
-            # Find a tmp_dir if we're supposed to use one.
-            tmp_dir = None
-            if spack.use_tmp_stage:
-                tmp_dir = next((tmp for tmp in spack.tmp_dirs
-                                if can_access(tmp)), None)
-
-            if not tmp_dir:
-                # If we couldn't find a tmp dir or if we're not using tmp
-                # stages, create the stage directly in spack.stage_path.
-                spack.mkdirp(self.path)
+        # Create the top-level stage directory
+        spack.mkdirp(spack.stage_path)
+        self._cleanup_dead_links()
+
+        # If this is a named stage, then construct a named path.
+        if self.name is not None:
+            self.path = new_path(spack.stage_path, self.name)
+
+        # If this is a temporary stage, them make the temp directory
+        tmp_dir = None
+        if self.tmp_root:
+            if self.name is None:
+                # Unnamed tmp root.  Link the path in
+                tmp_dir = tempfile.mkdtemp('', STAGE_PREFIX, self.tmp_root)
+                self.name = os.path.basename(tmp_dir)
+                self.path = new_path(spack.stage_path, self.name)
+                if self._need_to_create_path():
+                    os.symlink(tmp_dir, self.path)
 
             else:
-                # Otherwise we found a tmp_dir, so create the stage there
-                # and link it back to the prefix.
-                username = getpass.getuser()
-                if username:
-                    tmp_dir = spack.new_path(tmp_dir, username)
-                spack.mkdirp(tmp_dir)
-                tmp_dir = tempfile.mkdtemp('.stage', 'spack-stage-', tmp_dir)
+                if self._need_to_create_path():
+                    tmp_dir = tempfile.mkdtemp('', STAGE_PREFIX, self.tmp_root)
+                    os.symlink(tmp_dir, self.path)
 
-                os.symlink(tmp_dir, self.path)
+        # if we're not using a tmp dir, create the stage directly in the
+        # stage dir, rather than linking to it.
+        else:
+            if self.name is None:
+                self.path = tempfile.mkdtemp('', STAGE_PREFIX, spack.stage_path)
+                self.name = os.path.basename(self.path)
+            else:
+                if self._need_to_create_path():
+                    mkdirp(self.path)
 
         # Make sure we can actually do something with the stage we made.
         ensure_access(self.path)
@@ -187,7 +230,7 @@ class Stage(object):
         if not self.archive_file:
             tty.die("Attempt to expand archive before fetching.")
 
-        decompress = spack.decompressor_for(self.archive_file)
+        decompress = decompressor_for(self.archive_file)
         decompress(self.archive_file)
 
 
@@ -252,3 +295,22 @@ def purge():
         for stage_dir in os.listdir(spack.stage_path):
             stage_path = spack.new_path(spack.stage_path, stage_dir)
             remove_linked_tree(stage_path)
+
+
+def find_tmp_root():
+    if spack.use_tmp_stage:
+        for tmp in spack.tmp_dirs:
+            try:
+                mkdirp(expand_user(tmp))
+                return tmp
+            except OSError:
+                continue
+    return None
+
+
+class FailedDownloadError(serr.SpackError):
+    """Raised wen a download fails."""
+    def __init__(self, url):
+        super(FailedDownloadError, self).__init__(
+            "Failed to fetch file from URL: " + url)
+        self.url = url
diff --git a/lib/spack/spack/test/stage.py b/lib/spack/spack/test/stage.py
new file mode 100644
index 0000000000..19c0ed2fb3
--- /dev/null
+++ b/lib/spack/spack/test/stage.py
@@ -0,0 +1,242 @@
+"""\
+Test that the Stage class works correctly.
+"""
+import unittest
+import shutil
+import os
+import getpass
+from contextlib import *
+
+import spack
+from spack.stage import Stage
+from spack.util.filesystem import *
+from spack.util.executable import which
+
+test_files_dir = new_path(spack.stage_path, '.test')
+test_tmp_path  = new_path(test_files_dir, 'tmp')
+
+archive_dir      = 'test-files'
+archive_name     = archive_dir + '.tar.gz'
+archive_dir_path = new_path(test_files_dir, archive_dir)
+archive_url      = 'file://' + new_path(test_files_dir, archive_name)
+readme_name      = 'README.txt'
+test_readme      = new_path(archive_dir_path, readme_name)
+readme_text      = "hello world!\n"
+
+stage_name = 'spack-test-stage'
+
+
+class with_tmp(object):
+    """Decorator that executes a function with or without spack set
+       to use a temp dir."""
+    def __init__(self, use_tmp):
+        self.use_tmp = use_tmp
+
+    def __call__(self, fun):
+        use_tmp = self.use_tmp
+        def new_test_function(self):
+            old_tmp = spack.use_tmp_stage
+            spack.use_tmp_stage = use_tmp
+            fun(self)
+            spack.use_tmp_stage = old_tmp
+        return new_test_function
+
+
+class StageTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        """This sets up a mock archive to fetch, and a mock temp space for use
+           by the Stage class.  It doesn't actually create the Stage -- that
+           is done by individual tests.
+        """
+        if os.path.exists(test_files_dir):
+            shutil.rmtree(test_files_dir)
+
+        mkdirp(test_files_dir)
+        mkdirp(archive_dir_path)
+        mkdirp(test_tmp_path)
+
+        with closing(open(test_readme, 'w')) as readme:
+            readme.write(readme_text)
+
+        with working_dir(test_files_dir):
+            tar = which('tar')
+            tar('czf', archive_name, archive_dir)
+
+        # Make spack use the test environment for tmp stuff.
+        cls.old_tmp_dirs = spack.tmp_dirs
+        spack.tmp_dirs = [test_tmp_path]
+
+
+    @classmethod
+    def tearDownClass(cls):
+        """Blows away the test environment directory."""
+        shutil.rmtree(test_files_dir)
+
+        # restore spack's original tmp environment
+        spack.tmp_dirs = cls.old_tmp_dirs
+
+
+    def get_stage_path(self, stage, stage_name):
+        """Figure out based on a stage and an intended name where it should
+           be living.  This depends on whether it's named or not.
+        """
+        if stage_name:
+            # If it is a named stage, we know where the stage should be
+            stage_path = new_path(spack.stage_path, stage_name)
+        else:
+            # If it's unnamed, ensure that we ran mkdtemp in the right spot.
+            stage_path = stage.path
+            self.assertIsNotNone(stage_path)
+            self.assertEqual(
+                os.path.commonprefix((stage_path, spack.stage_path)),
+                spack.stage_path)
+        return stage_path
+
+
+    def check_setup(self, stage, stage_name):
+        """Figure out whether a stage was set up correctly."""
+        stage_path = self.get_stage_path(stage, stage_name)
+        self.assertTrue(os.path.isdir(stage_path))
+
+        if spack.use_tmp_stage:
+            # Make sure everything was created and linked correctly for
+            # a tmp stage.
+            self.assertTrue(os.path.islink(stage_path))
+
+            target = os.path.realpath(stage_path)
+            self.assertTrue(os.path.isdir(target))
+            self.assertFalse(os.path.islink(target))
+            self.assertEqual(
+                os.path.commonprefix((target, test_tmp_path)),
+                test_tmp_path)
+
+        else:
+            # Make sure the stage path is NOT a link for a non-tmp stage
+            self.assertFalse(os.path.islink(stage_path))
+
+
+    def check_fetch(self, stage, stage_name):
+        stage_path = self.get_stage_path(stage, stage_name)
+        self.assertTrue(archive_name in os.listdir(stage_path))
+        self.assertEqual(new_path(stage_path, archive_name),
+                         stage.archive_file)
+
+
+    def check_expand_archive(self, stage, stage_name):
+        stage_path = self.get_stage_path(stage, stage_name)
+        self.assertTrue(archive_name in os.listdir(stage_path))
+        self.assertTrue(archive_dir in os.listdir(stage_path))
+
+        readme = new_path(stage_path, archive_dir, readme_name)
+        self.assertTrue(os.path.isfile(readme))
+
+        with closing(open(readme)) as file:
+            self.assertEqual(readme_text, file.read())
+
+
+    def check_chdir(self, stage, stage_name):
+        stage_path = self.get_stage_path(stage, stage_name)
+        self.assertEqual(os.path.realpath(stage_path), os.getcwd())
+
+
+    def check_chdir_to_archive(self, stage, stage_name):
+        stage_path = self.get_stage_path(stage, stage_name)
+        self.assertEqual(
+            new_path(os.path.realpath(stage_path), archive_dir),
+            os.getcwd())
+
+
+    def check_destroy(self, stage, stage_name):
+        """Figure out whether a stage was destroyed correctly."""
+        stage_path = self.get_stage_path(stage, stage_name)
+
+        # check that the stage dir/link was removed.
+        self.assertFalse(os.path.exists(stage_path))
+
+        # tmp stage needs to remove tmp dir too.
+        if spack.use_tmp_stage:
+            target = os.path.realpath(stage_path)
+            self.assertFalse(os.path.exists(target))
+
+
+    def checkSetupAndDestroy(self, stage_name=None):
+        stage = Stage(archive_url, stage_name)
+        stage.setup()
+        self.check_setup(stage, stage_name)
+
+        stage.destroy()
+        self.check_destroy(stage, stage_name)
+
+
+    @with_tmp(True)
+    def test_setup_and_destroy_name_with_tmp(self):
+        self.checkSetupAndDestroy(stage_name)
+
+
+    @with_tmp(False)
+    def test_setup_and_destroy_name_without_tmp(self):
+        self.checkSetupAndDestroy(stage_name)
+
+
+    @with_tmp(True)
+    def test_setup_and_destroy_no_name_with_tmp(self):
+        self.checkSetupAndDestroy(None)
+
+
+    @with_tmp(False)
+    def test_setup_and_destroy_no_name_without_tmp(self):
+        self.checkSetupAndDestroy(None)
+
+
+    def test_chdir(self):
+        stage = Stage(archive_url, stage_name)
+
+        stage.chdir()
+        self.check_setup(stage, stage_name)
+        self.check_chdir(stage, stage_name)
+
+        stage.destroy()
+        self.check_destroy(stage, stage_name)
+
+
+    def test_fetch(self):
+        stage = Stage(archive_url, stage_name)
+
+        stage.fetch()
+        self.check_setup(stage, stage_name)
+        self.check_chdir(stage, stage_name)
+        self.check_fetch(stage, stage_name)
+
+        stage.destroy()
+        self.check_destroy(stage, stage_name)
+
+
+    def test_expand_archive(self):
+        stage = Stage(archive_url, stage_name)
+
+        stage.fetch()
+        self.check_setup(stage, stage_name)
+        self.check_fetch(stage, stage_name)
+
+        stage.expand_archive()
+        self.check_expand_archive(stage, stage_name)
+
+        stage.destroy()
+        self.check_destroy(stage, stage_name)
+
+
+    def test_zexpand_archive(self):
+        stage = Stage(archive_url, stage_name)
+
+        stage.fetch()
+        self.check_setup(stage, stage_name)
+        self.check_fetch(stage, stage_name)
+
+        stage.expand_archive()
+        stage.chdir_to_archive()
+        self.check_expand_archive(stage, stage_name)
+        self.check_chdir_to_archive(stage, stage_name)
+
+        stage.destroy()
+        self.check_destroy(stage, stage_name)
diff --git a/lib/spack/spack/util/compression.py b/lib/spack/spack/util/compression.py
index b0dc0241e3..c4ac256826 100644
--- a/lib/spack/spack/util/compression.py
+++ b/lib/spack/spack/util/compression.py
@@ -1,4 +1,5 @@
 from itertools import product
+from spack.util.executable import which
 
 # Supported archvie extensions.
 PRE_EXTS = ["tar"]
diff --git a/lib/spack/spack/util/filesystem.py b/lib/spack/spack/util/filesystem.py
index 8188946ccb..32f99af533 100644
--- a/lib/spack/spack/util/filesystem.py
+++ b/lib/spack/spack/util/filesystem.py
@@ -2,17 +2,29 @@ import os
 import re
 import shutil
 import errno
+import getpass
 from contextlib import contextmanager, closing
 
 import spack.tty as tty
 from spack.util.compression import ALLOWED_ARCHIVE_TYPES
 
+
 def install(src, dest):
     """Manually install a file to a particular location."""
     tty.info("Installing %s to %s" % (src, dest))
     shutil.copy(src, dest)
 
 
+def expand_user(path):
+    """Find instances of '%u' in a path and replace with the current user's
+       username."""
+    username = getpass.getuser()
+    if not username and '%u' in path:
+        tty.die("Couldn't get username to complete path '%s'" % path)
+
+    return path.replace('%u', username)
+
+
 @contextmanager
 def working_dir(dirname):
     orig_dir = os.getcwd()