Seperate snapshoting from get_device in Disk class

[snf-image-creator] / image_creator / disk.py

diff --git a/image_creator/disk.py b/image_creator/disk.py
index 809a1a3..ef53799 100644 (file)
--- a/image_creator/disk.py
+++ b/image_creator/disk.py
@@ -1,7 +1,38 @@
-#!/usr/bin/env python
+# Copyright 2012 GRNET S.A. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or
+# without modification, are permitted provided that the following
+# conditions are met:
+#
+#   1. Redistributions of source code must retain the above
+#      copyright notice, this list of conditions and the following
+#      disclaimer.
+#
+#   2. Redistributions in binary form must reproduce the above
+#      copyright notice, this list of conditions and the following
+#      disclaimer in the documentation and/or other materials
+#      provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
+# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
+# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
+# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+#
+# The views and conclusions contained in the software and
+# documentation are those of the authors and should not be
+# interpreted as representing official policies, either expressed
+# or implied, of GRNET S.A.

 
 from image_creator.util import get_command
 
 from image_creator.util import get_command

-from clint.textui import progress
+from image_creator.util import warn, progress, success, output, FatalError

 
 import stat
 import os
 
 import stat
 import os

@@ -10,6 +41,8 @@ import uuid
 import re
 import sys
 import guestfs
 import re
 import sys
 import guestfs

+import time
+from sendfile import sendfile

 
 
 class DiskError(Exception):
 
 
 class DiskError(Exception):

@@ -41,7 +74,7 @@ class Disk(object):
 
     def _losetup(self, fname):
         loop = losetup('-f', '--show', fname)
 
     def _losetup(self, fname):
         loop = losetup('-f', '--show', fname)

-        loop = loop.strip() # remove the new-line char
+        loop = loop.strip()  # remove the new-line char

         self._add_cleanup(losetup, '-d', loop)
         return loop
 
         self._add_cleanup(losetup, '-d', loop)
         return loop
 

@@ -60,27 +93,35 @@ class Disk(object):
             job, args = self._cleanup_jobs.pop()
             job(*args)
 
             job, args = self._cleanup_jobs.pop()
             job(*args)
 

-    def get_device(self):
-        """Returns a newly created DiskDevice instance.
-
-        This instance is a snapshot of the original source media of
-        the Disk instance.
+    def snapshot(self):
+        """Creates a snapshot of the original source media of the Disk
+        instance.

         """
         """

+
+        output("Examining source media `%s'..." % self.source, False)

         sourcedev = self.source
         mode = os.stat(self.source).st_mode
         if stat.S_ISDIR(mode):
         sourcedev = self.source
         mode = os.stat(self.source).st_mode
         if stat.S_ISDIR(mode):

+            success('looks like a directory')

             return self._losetup(self._dir_to_disk())
         elif stat.S_ISREG(mode):
             return self._losetup(self._dir_to_disk())
         elif stat.S_ISREG(mode):

+            success('looks like an image file')

             sourcedev = self._losetup(self.source)
         elif not stat.S_ISBLK(mode):
             sourcedev = self._losetup(self.source)
         elif not stat.S_ISBLK(mode):

-            raise ValueError("Value for self.source is invalid")
+            raise ValueError("Invalid media source. Only block devices, "
+                            "regular files and directories are supported.")
+        else:
+            success('looks like a block device')

 
         # Take a snapshot and return it to the user
 
         # Take a snapshot and return it to the user

+        output("Snapshotting media source...", False)

         size = blockdev('--getsize', sourcedev)
         cowfd, cow = tempfile.mkstemp()
         size = blockdev('--getsize', sourcedev)
         cowfd, cow = tempfile.mkstemp()

+        os.close(cowfd)

         self._add_cleanup(os.unlink, cow)
         # Create 1G cow sparse file
         self._add_cleanup(os.unlink, cow)
         # Create 1G cow sparse file

-        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
+        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
+                                        'seek=%d' % (1024 * 1024))

         cowdev = self._losetup(cow)
 
         snapshot = uuid.uuid4().hex
         cowdev = self._losetup(cow)
 
         snapshot = uuid.uuid4().hex

@@ -90,11 +131,22 @@ class Disk(object):
                                         (int(size), sourcedev, cowdev))
             dmsetup('create', snapshot, table)
             self._add_cleanup(dmsetup, 'remove', snapshot)
                                         (int(size), sourcedev, cowdev))
             dmsetup('create', snapshot, table)
             self._add_cleanup(dmsetup, 'remove', snapshot)

+            # Sometimes dmsetup remove fails with Device or resource busy,
+            # although everything is cleaned up and the snapshot is not
+            # used by anyone. Add a 2 seconds delay to be on the safe side.
+            self._add_cleanup(time.sleep, 2)
+

         finally:
             os.unlink(table)
         finally:
             os.unlink(table)

+        success('done')
+        return "/dev/mapper/%s" % snapshot

 
 

-        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
+    def get_device(self, media):
+        """Returns a newly created DiskDevice instance."""
+
+        new_device = DiskDevice(media)

         self._devices.append(new_device)
         self._devices.append(new_device)

+        new_device.enable()

         return new_device
 
     def destroy_device(self, device):
         return new_device
 
     def destroy_device(self, device):

@@ -105,15 +157,6 @@ class Disk(object):
         device.destroy()
 
 
         device.destroy()
 
 

-def progress_generator(total):
-    position = 0;
-    for i in progress.bar(range(total)):
-        if i < position:
-            continue
-        position = yield
-    yield #suppress the StopIteration exception
-
-

 class DiskDevice(object):
     """This class represents a block device hosting an Operating System
     as created by the device-mapper.
 class DiskDevice(object):
     """This class represents a block device hosting an Operating System
     as created by the device-mapper.

@@ -121,52 +164,67 @@ class DiskDevice(object):
 
     def __init__(self, device, bootable=True):
         """Create a new DiskDevice."""
 
     def __init__(self, device, bootable=True):
         """Create a new DiskDevice."""

+

         self.device = device
         self.bootable = bootable
         self.progress_bar = None
 
         self.g = guestfs.GuestFS()
         self.device = device
         self.bootable = bootable
         self.progress_bar = None
 
         self.g = guestfs.GuestFS()

-        self.g.add_drive_opts(device, readonly=0)
+        self.g.add_drive_opts(self.device, readonly=0)

 
         #self.g.set_trace(1)
         #self.g.set_verbose(1)
 
 
         #self.g.set_trace(1)
         #self.g.set_verbose(1)
 

-        eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS)
+        self.guestfs_enabled = False
+
+    def enable(self):
+        """Enable a newly created DiskDevice"""
+        self.progressbar = progress("Launching helper VM: ")
+        self.progressbar.next()
+        eh = self.g.set_event_callback(self.progress_callback,
+                                                    guestfs.EVENT_PROGRESS)

         self.g.launch()
         self.g.launch()

+        self.guestfs_enabled = True

         self.g.delete_event_callback(eh)
         self.g.delete_event_callback(eh)

-        
+        if self.progressbar is not None:
+            self.progressbar.send(100)
+            self.progressbar = None
+
+        output('Inspecting Operating System...', False)

         roots = self.g.inspect_os()
         if len(roots) == 0:
         roots = self.g.inspect_os()
         if len(roots) == 0:

-            raise DiskError("No operating system found")
+            raise FatalError("No operating system found")

         if len(roots) > 1:
         if len(roots) > 1:

-            raise DiskError("Multiple operating systems found")
-
+            raise FatalError("Multiple operating systems found."
+                            "We only support images with one filesystem.")

         self.root = roots[0]
         self.ostype = self.g.inspect_get_type(self.root)
         self.distro = self.g.inspect_get_distro(self.root)
         self.root = roots[0]
         self.ostype = self.g.inspect_get_type(self.root)
         self.distro = self.g.inspect_get_distro(self.root)

+        success('found a %s system' % self.distro)

 
     def destroy(self):
         """Destroy this DiskDevice instance."""
 
     def destroy(self):
         """Destroy this DiskDevice instance."""

-        self.g.umount_all()
-        self.g.sync()
-        # Close the guestfs handler
+
+        if self.guestfs_enabled:
+            self.g.umount_all()
+            self.g.sync()
+
+        # Close the guestfs handler if open

         self.g.close()
 
     def progress_callback(self, ev, eh, buf, array):
         position = array[2]
         total = array[3]
         self.g.close()
 
     def progress_callback(self, ev, eh, buf, array):
         position = array[2]
         total = array[3]

-        
-        if self.progress_bar is None:
-            self.progress_bar = progress_generator(total)
-            self.progress_bar.next()

 
 

-        self.progress_bar.send(position)
+        self.progressbar.send((position * 100) // total)

 
         if position == total:
 
         if position == total:

-            self.progress_bar = None
+            self.progressbar = None

 
     def mount(self):
         """Mount all disk partitions in a correct order."""
 
     def mount(self):
         """Mount all disk partitions in a correct order."""

+
+        output("Mounting image...", False)

         mps = self.g.inspect_get_mountpoints(self.root)
 
         # Sort the keys to mount the fs in a correct order.
         mps = self.g.inspect_get_mountpoints(self.root)
 
         # Sort the keys to mount the fs in a correct order.

@@ -183,7 +241,8 @@ class DiskDevice(object):
             try:
                 self.g.mount(dev, mp)
             except RuntimeError as msg:
             try:
                 self.g.mount(dev, mp)
             except RuntimeError as msg:

-                print "%s (ignored)" % msg
+                warn("%s (ignored)" % msg)
+        success("done")

 
     def umount(self):
         """Umount all mounted filesystems."""
 
     def umount(self):
         """Umount all mounted filesystems."""

@@ -196,29 +255,34 @@ class DiskDevice(object):
         disk and then updating the partition table. The new disk size
         (in bytes) is returned.
         """
         disk and then updating the partition table. The new disk size
         (in bytes) is returned.
         """

+        output("Shrinking image (this may take a while)...", False)
+

         dev = self.g.part_to_dev(self.root)
         parttype = self.g.part_get_parttype(dev)
         if parttype != 'msdos':
         dev = self.g.part_to_dev(self.root)
         parttype = self.g.part_get_parttype(dev)
         if parttype != 'msdos':

-            raise DiskError("You have a %s partition table. "
+            raise FatalError("You have a %s partition table. "

                 "Only msdos partitions are supported" % parttype)
 
         last_partition = self.g.part_list(dev)[-1]
 
         if last_partition['part_num'] > 4:
                 "Only msdos partitions are supported" % parttype)
 
         last_partition = self.g.part_list(dev)[-1]
 
         if last_partition['part_num'] > 4:

-            raise DiskError("This disk contains logical partitions. "
+            raise FatalError("This disk contains logical partitions. "

                 "Only primary partitions are supported.")
 
         part_dev = "%s%d" % (dev, last_partition['part_num'])
         fs_type = self.g.vfs_type(part_dev)
         if not re.match("ext[234]", fs_type):
                 "Only primary partitions are supported.")
 
         part_dev = "%s%d" % (dev, last_partition['part_num'])
         fs_type = self.g.vfs_type(part_dev)
         if not re.match("ext[234]", fs_type):

-            print "Warning: Don't know how to resize %s partitions." % vfs_type
+            warn("Don't know how to resize %s partitions." % vfs_type)

             return
 
         self.g.e2fsck_f(part_dev)
         self.g.resize2fs_M(part_dev)
             return
 
         self.g.e2fsck_f(part_dev)
         self.g.resize2fs_M(part_dev)

-        output = self.g.tune2fs_l(part_dev)
-        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
-        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
+
+        out = self.g.tune2fs_l(part_dev)
+        block_size = int(
+            filter(lambda x: x[0] == 'Block size', out)[0][1])
+        block_cnt = int(
+            filter(lambda x: x[0] == 'Block count', out)[0][1])

 
         sector_size = self.g.blockdev_getss(dev)
 
 
         sector_size = self.g.blockdev_getss(dev)
 

@@ -228,7 +292,10 @@ class DiskDevice(object):
         self.g.part_del(dev, last_partition['part_num'])
         self.g.part_add(dev, 'p', start, end)
 
         self.g.part_del(dev, last_partition['part_num'])
         self.g.part_add(dev, 'p', start, end)
 

-        return (end + 1) * sector_size
+        new_size = (end + 1) * sector_size
+        success("new image size is %dMB" %
+                            ((new_size + 2 ** 20 - 1) // 2 ** 20))
+        return new_size

 
     def size(self):
         """Returns the "payload" size of the device.
 
     def size(self):
         """Returns the "payload" size of the device.

@@ -239,6 +306,39 @@ class DiskDevice(object):
         dev = self.g.part_to_dev(self.root)
         last = self.g.part_list(dev)[-1]
 
         dev = self.g.part_to_dev(self.root)
         last = self.g.part_list(dev)[-1]
 

-        return last['part_end']
+        return last['part_end'] + 1
+
+    def dump(self, outfile):
+        """Dumps the content of device into a file.
+
+        This method will only dump the actual payload, found by reading the
+        partition table. Empty space in the end of the device will be ignored.
+        """
+        blocksize = 2 ** 22  # 4MB
+        size = self.size()
+        progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
+        progressbar = progress("Dumping image file: ", progress_size)
+
+        source = open(self.device, "r")
+        try:
+            dest = open(outfile, "w")
+            try:
+                left = size
+                offset = 0
+                progressbar.next()
+                while left > 0:
+                    length = min(left, blocksize)
+                    sent = sendfile(dest.fileno(), source.fileno(), offset,
+                                                                        length)
+                    offset += sent
+                    left -= sent
+                    for i in range((length + 2 ** 20 - 1) // 2 ** 20):
+                        progressbar.next()
+            finally:
+                dest.close()
+        finally:
+            source.close()
+
+        success('Image file %s was successfully created' % outfile)

 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :