Commits

Albert Hopkins committed 120d494

Refactor: make get_image() into a proper context manager.
Various other fixes

  • Participants
  • Parent commits b8694a5
  • Tags 0.4.9.2

Comments (0)

Files changed (3)

 
 from contextlib import contextmanager
 import os
-import re
 import select
 import SocketServer
 from subprocess import Popen, PIPE
         status = sudo_popen(['fuser', '-a', devnode], stderr=NULL).wait()
         if status != 0:
             args = ['qemu-nbd', '--connect=%s' % devnode, filename]
-            sudo_popen(args, stderr=NULL, stdout=NULL, stdin=NULL, cwd=dirname)
-            sleep(WAIT_TIME)
+            sudo_popen(args, stderr=NULL, stdout=NULL, stdin=NULL,
+                    cwd=dirname).wait()
+            sleep(1) # it never works instantly :(
             yield devnode
             break
 
         dev_num = (dev_num + 1 if dev_num < NUM_DEVS else 0)
+    sudo_popen(('qemu-nbd', '-d', devnode), stdout=NULL).wait()
 
 def make_snapshot_image(directory, base_image, snapshot_image):
     """Create a qcow2 snapshot image based on «base_image» in «directory».
     status = Popen(args, stdout=NULL, cwd=directory).wait()
     return status
 
-def mount_qcow(image, mount_point, partition=1):
-    "Map the qcow image «image» at «mount_point»"
+def mount_qcow(devnode, mount_point, partition=1):
+    """Map the nbd device «devnode» at «mount_point».  If «mount_point» does
+    not exist, attempt to create it."""
 
-    with get_device(image) as dev:
+    if not os.path.isdir(mount_point):
         os.makedirs(mount_point)
-        for dummy in range(ATTEMPTS):
-            # we attempt this a few times because it sometimes fails
-            # the first time under high load
+    for dummy in range(ATTEMPTS):
+        # we attempt this a few times because it sometimes fails
+        # the first time under high load
 
-            sleep(1) # we wait first
-            if partition == 0:
-                status = sudo_popen(['mount', dev, mount_point]).wait()
-            else:
-                status = sudo_popen(['mount', '%sp%s' % (dev, partition),
-                    mount_point]).wait()
-            if status == 0:
-                break
+        sleep(1) # we wait first
+        if partition == 0:
+            status = sudo_popen(['mount', '-o', 'noatime', devnode,
+                mount_point]).wait()
         else:
-            raise OSError, ('Failed to mount %s to %s after %s attempts'
-                    % (dev, mount_point, ATTEMPTS))
-
-def get_device_for_mountpoint(mount_point):
-    """Returns what device node is mounted on «mount_point» or None if not
-    mounted."""
-
-    regex = re.compile(MTAB % mount_point)
-    with open('/proc/mounts') as f_object:
-        for line in f_object:
-            match_obj = regex.match(line)
-            if match_obj:
-                dev = match_obj.group('dev')
-                return dev
+            status = sudo_popen(['mount', '-o', 'noatime', '%sp%s' %
+                (devnode, partition), mount_point]).wait()
+        if status == 0:
+            break
+    else:
+        raise OSError, ('Failed to mount %s to %s after %s attempts'
+                % (devnode, mount_point, ATTEMPTS))
 
 def unmount_qcow(mount_point):
-    """unmounts qcow device mounted at «mount_point» and disconnects device
-    from the nbd"""
-
-    dev = get_device_for_mountpoint(mount_point)
-    if not dev:
-        return
+    """unmounts qcow device mounted at «mount_point». Removes the directory"""
 
     for dummy in range(ATTEMPTS):
         status = sudo_popen(['umount', mount_point], stderr=NULL).wait()
         if status == 0:
-            sudo_popen(['qemu-nbd', '-d', dev], stdout=NULL).wait()
             os.removedirs(mount_point)
             break
         sleep(WAIT_TIME)
     """
     xml_start_pos = 92
     xml_end_pos = 4096
-    null = '\x00'
 
-    state = open(state_file, 'r+b')
-    state.seek(xml_start_pos)
+    with open(state_file, 'r+b') as state:
+        state.seek(xml_start_pos)
 
-    # up to xml_end_pos contains the dom xml, right padded with null bytes
-    xml = state.read(xml_end_pos - xml_start_pos).rstrip(null)
-    domtree = ElementTree.XML(xml)
+        # up to xml_end_pos contains the dom xml, right padded with null bytes
+        xml = state.read(xml_end_pos - xml_start_pos).rstrip('\x00')
+        domtree = ElementTree.XML(xml)
 
-    disk = domtree.find('devices/disk/source')
-    if disk is not None:
-        orig_file = disk.get('file')
-        basename = os.path.basename(orig_file)
-        new_file = '%s/%s' % (new_image_dir, basename)
-        disk.set('file', new_file)
-
-    for path in ('os/kernel', 'os/initrd'):
-        tag = domtree.find(path)
+        tag = domtree.find('devices/disk/source')
         if tag is not None:
-            filename = tag.text
+            filename = tag.get('file')
             basename = os.path.basename(filename)
             new_file = '%s/%s' % (new_image_dir, basename)
-            tag.text = new_file
+            tag.set('file', new_file)
 
-    # we append a newline because that's what the original had
-    xml = ElementTree.tostring(domtree) + '\n'
+        for path in ('os/kernel', 'os/initrd'):
+            tag = domtree.find(path)
+            if tag is not None:
+                filename = tag.text
+                basename = os.path.basename(filename)
+                new_file = '%s/%s' % (new_image_dir, basename)
+                tag.text = new_file
 
-    state.seek(xml_start_pos)
-    state.write(xml)
+        # we append a newline because that's what the original had
+        xml = ElementTree.tostring(domtree) + '\n'
 
-    # fill up to xml_end_pos with nulls
-    pos = state.tell()
-    state.write(null * (xml_end_pos - pos))
+        state.seek(xml_start_pos)
+        state.write(xml)
 
-    state.close()
+        # fill up to xml_end_pos with nulls
+        pos = state.tell()
+        state.write('\x00' * (xml_end_pos - pos))
 
 def which(progname, path=None, test=None):
     """Like the 'which' command, but works on any type of file/dir, not
 from subprocess import call, PIPE, Popen
 import sys
 import tempfile
-import time
 import xml.etree.ElementTree as ElementTree
 
 import libvirt
 
-from helpers import get_device, sudo_popen
+from helpers import get_device, sudo_popen, mount_qcow, unmount_qcow
 import settings
 
-__version__ = '0.4.9.1'
+__version__ = '0.4.9.2'
 __revision__ = '$Revision$'
 
 DIRS = ('.hemp', '.hemp/images',)
     call(('qemu-img', 'create', '-f', 'qcow2', tmpfile, str(project_size)),
             stdout=NULL)
     with get_device(tmpfile) as devnode:
-        time.sleep(3)
-        sudo_popen(('/sbin/mkfs', '-m', '0', '-t', 'ext4', '-L', name,
+        status = sudo_popen(('/sbin/mkfs', '-m', '0', '-t', 'ext4', '-L', name,
             '-q', '-O', '^has_journal', devnode)).wait()
+        if status != 0:
+            return
         tmpdir = tempfile.mkdtemp()
-        sudo_popen(('mount', '-t', 'ext4', '-o', 'noatime',  devnode,
-            tmpdir)).wait()
+        mount_qcow(devnode, tmpdir, 0)
         # copy files into the temporary mount point
         sudo_popen(('/bin/cp', '-a', args.path + '/.', tmpdir)).wait()
         sudo_popen(('/bin/chmod', '0700', tmpdir)).wait()
         # umount it
-        sudo_popen(('umount', tmpdir)).wait()
-        os.rmdir(tmpdir)
+        unmount_qcow(tmpdir)
         sudo_popen(('qemu-nbd', '-d', devnode), stdout=NULL).wait()
 
     # now we re-create the .qcow and compress it
     """Unpack a "live" cloud"""
     filename = os.path.realpath(args.filename)
     with get_device(filename) as devnode:
-        time.sleep(3)
         popen = sudo_popen(('/sbin/e2label', devnode), stdout=PIPE)
         name = popen.communicate()[0].strip()
 
             name)
         popen.wait()
         tmpdir = tempfile.mkdtemp(prefix='hemp.%s.' % name)
-        sudo_popen(('mount', '-t', 'ext4', '-o', 'noatime', devnode,
-            tmpdir)).wait()
+        mount_qcow(devnode, tmpdir, 0)
         sudo_popen(('chown', '-R', 'nobody:nobody', tmpdir)).wait()
         rcfilename = tempfile.mkstemp(prefix='hemp.%s.' % name)[1]
-        rcfile = open(rcfilename, 'w')
-        rcfile.write('cd "%s"\n' % tmpdir)
-        rcfile.write('PROMPT_COMMAND="fab status"\n')
-        rcfile.write('PS1="%s"\n' % prompt)
-        rcfile.write('export PYTHONDONTWRITEBYTECODE=1\n')
-        rcfile.write('clear\n')
-        rcfile.write("echo -e ' \\e[0;32m%s'\n" % ('-' * 78))
-        rcfile.write("echo '%s'\n" % 'Hemp Live Cloud'.center(80))
-        rcfile.write("echo ' %s'\n" % ('-' * 78))
-        rcfile.write("echo -en "
-            "'  \\e[0;31mYou are root inside a hemp shell.'\n")
-        rcfile.write("echo ' Remember to \"fab shutdown\" before exiting.'\n")
-        rcfile.write("echo -e '\\e[0m'\n")
-        if args.run:
-            rcfile.write('fab go\n')
-        rcfile.close()
+        with open(rcfilename, 'w') as rcfile:
+            rwrite = rcfile.write
+            rwrite('cd "%s"\n' % tmpdir)
+            rwrite('PROMPT_COMMAND="fab status"\n')
+            rwrite('PS1="%s"\n' % prompt)
+            rwrite('export PYTHONDONTWRITEBYTECODE=1\n')
+            rwrite('clear\n')
+            rwrite("echo -e ' \\e[0;32m%s'\n" % ('-' * 78))
+            rwrite("echo '%s'\n" % 'Hemp Live Cloud'.center(80))
+            rwrite("echo ' %s'\n" % ('-' * 78))
+            rwrite("echo -en "
+                "'  \\e[0;31mYou are root inside a hemp shell.'\n")
+            rwrite("echo ' Remember to \"fab shutdown\" before exiting.'\n")
+            rwrite("echo -e '\\e[0m'\n")
+            if args.run:
+                rwrite('fab go\n')
 
         # run a restricted shell into the hemp root. wait() for the shell to
         # exit
         # umount and disconnect
         # since the filesystem may be busy (e.g. hosts are running), we
         # need to do this in a loop
-        while True:
-            status = sudo_popen(('umount', tmpdir)).wait()
-            if status == 0:
-                break
-            sys.stderr.write('Error closing cloud file, retrying in 10 '
-                'seconds.\nDid you forget to shutdown?.\n')
-            time.sleep(10)
-        time.sleep(3)
-        sudo_popen(('qemu-nbd', '-d', devnode), stdout=NULL).wait()
-
-        # cleanup tmp file/dir
-        os.rmdir(tmpdir)
+        unmount_qcow(tmpdir)
         os.unlink(rcfilename)
 
 def pack(args):
     keydest = getattr(config, 'KEY_DEST', None)
 
     CLOUD.helpers.make_snapshot_image(image_dir, BASE_IMG, image)
-    CLOUD.helpers.mount_qcow('%s/%s' % (image_dir, image), mount_point)
-    CLOUD.helpers.set_image_to_hostname(mount_point, env.host)
-    CLOUD.helpers.remove_persistent_net_rules(mount_point)
-    if keyfile:
-        CLOUD.helpers.install_keyfile(os.path.expanduser(keyfile),
-                mount_point, keydest)
-    CLOUD.helpers.unmount_qcow(mount_point)
+    with CLOUD.helpers.get_device('%s/%s' % (image_dir, image)) as devnode:
+        CLOUD.helpers.mount_qcow(devnode, mount_point)
+        CLOUD.helpers.set_image_to_hostname(mount_point, env.host)
+        CLOUD.helpers.remove_persistent_net_rules(mount_point)
+        if keyfile:
+            CLOUD.helpers.install_keyfile(os.path.expanduser(keyfile),
+                    mount_point, keydest)
+        CLOUD.helpers.unmount_qcow(mount_point)
 
 def rm_netrules():
     """[H] Remove the persitent-net udev rules file"""
     mount_point = '%s/.hemp/mnt/%s' % (DIRNAME, env.host)
     image = '%s.qcow2' % env.host
     image_dir = '%s/.hemp/images' % DIRNAME
-    CLOUD.helpers.mount_qcow('%s/%s' % (image_dir, image), mount_point)
-    CLOUD.helpers.remove_persistent_net_rules(mount_point)
-    CLOUD.helpers.unmount_qcow(mount_point)
+    with CLOUD.helpers.get_device('%s/%s' % (image_dir, image)) as devnode:
+        CLOUD.helpers.mount_qcow(devnode, mount_point)
+        CLOUD.helpers.remove_persistent_net_rules(mount_point)
+        CLOUD.helpers.unmount_qcow(mount_point)
 
 def boot():
     """[H] Boot up the cloud"""