Merge "VDI-Streaming: support customizing compress level"

os_xenapi/client/image/__init__.py

@@ -22,7 +22,9 @@ def stream_to_vdis(context, session, instance, host_url, sr_ref, data):
     return handler.vdis
 
 
-def stream_from_vdis(context, session, instance, host_url, vdi_uuids):
+def stream_from_vdis(context, session, instance, host_url, vdi_uuids,
+                     compresslevel=vdi_handler.DEFAULT_COMPRESSLEVEL):
     handler = vdi_handler.GenerateImageStream(context, session, instance,
-                                              host_url, vdi_uuids)
+                                              host_url, vdi_uuids,
+                                              compresslevel)
     return handler.get_image_data()

os_xenapi/client/image/vdi_handler.py

@@ -26,6 +26,7 @@ LOG = logging.getLogger(__name__)
 
 
 CHUNK_SIZE = 4 * 1024 * 1024
+DEFAULT_COMPRESSLEVEL = 6
 
 
 class ImageStreamToVDIs(object):
@@ -129,19 +130,28 @@ class ImageStreamToVDIs(object):
 
 
 class GenerateImageStream(object):
-    def __init__(self, context, session, instance, host_url, vdi_uuids):
+    def __init__(self, context, session, instance, host_url, vdi_uuids,
+                 compresslevel=None):
         self.context = context
         self.session = session
         self.instance = instance
         self.host_url = host_url
         self.vdi_uuids = vdi_uuids
+        self.compresslevel = self._get_compress_level(compresslevel)
+
+    def _get_compress_level(self, compresslevel):
+        if compresslevel and compresslevel >= 1 and compresslevel <= 9:
+            return compresslevel
+        # if compresslevel is not set or not a valid value, just return
+        # the default value.
+        return DEFAULT_COMPRESSLEVEL
 
     def get_image_data(self):
         """This function will:
 
           1). export VDI as VHD stream;
           2). make gzipped tarball from the VHD stream;
-          3). read from the tarball stream.and return the iterable data.
+          3). read from the tarball stream and return the iterable data.
         """
 
         tarpipe_out, tarpipe_in = utils.create_pipe()
@@ -164,7 +174,7 @@ class GenerateImageStream(object):
     def start_image_stream_generator(self, tarpipe_in):
         tar_generator = VdisToTarStream(
             self.context, self.session, self.instance, self.host_url,
-            self.vdi_uuids, tarpipe_in)
+            self.vdi_uuids, tarpipe_in, self.compresslevel)
         try:
             tar_generator.start()
         finally:
@@ -173,19 +183,23 @@ class GenerateImageStream(object):
 
 class VdisToTarStream(object):
     def __init__(self, context, session, instance, host_url, vdi_uuids,
-                 tarpipe_in):
+                 tarpipe_in, compresslevel):
         self.context = context
         self.session = session
         self.instance = instance
         self.host_url = host_url
         self.vdi_uuids = vdi_uuids
         self.tarpipe_in = tarpipe_in
+        self.compresslevel = compresslevel
         self.conn = None
         self.task_ref = None
 
     def start(self):
         # Start thread to generate tgz and write tgz data into tarpipe_in.
-        with tarfile.open(fileobj=self.tarpipe_in, mode='w|gz') as tar_file:
+        with tarfile.TarFile.gzopen(
+            name=None, fileobj=self.tarpipe_in, mode='w',
+            compresslevel=self.compresslevel) as tar_file:
+
             # only need export the leaf vdi.
             vdi_uuid = self.vdi_uuids[0]
             vdi_ref = self.session.VDI.get_by_uuid(vdi_uuid)

os_xenapi/tests/client/image/test_init.py

@@ -37,8 +37,9 @@ class ImageTestCase(base.TestCase):
         mock_start.assert_called_once_with()
 
     @mock.patch.object(vdi_handler.GenerateImageStream, 'get_image_data')
-    def test_vdis_to_stream(self, mock_get):
+    def test_stream_from_vdis(self, mock_get):
         image.stream_from_vdis(self.context, self.session, self.instance,
-                               self.host_url, ['fake-uuid'])
+                               self.host_url, ['fake-uuid'],
+                               compresslevel=9)
 
         mock_get.assert_called_once_with()

os_xenapi/tests/client/image/test_vdi_handler.py

@@ -220,6 +220,69 @@ class GenerateImageStreamTestCase(base.TestCase):
         mock_spawn.assert_called_once_with(mock_start, mock_tarpipe_in)
         self.assertEqual(image_chunks, ['chunk1', 'chunk2'])
 
+    @mock.patch.object(vdi_handler, 'VdisToTarStream')
+    def test_start_stream_generator(self, mock_stream):
+        # Verify the specified compress level should be used,
+        # if a compresslevel specified in GenerateImageStream.
+        compr_level = 9
+        mock_stream_obj = mock.Mock()
+        mock_stream.return_value = mock_stream_obj
+        generator = vdi_handler.GenerateImageStream(
+            self.context, self.session, self.instance,
+            self.host_url, ['vdi_uuid'], compresslevel=compr_level)
+        fake_tarpipe_in = mock.Mock()
+
+        generator.start_image_stream_generator(fake_tarpipe_in)
+
+        mock_stream.assert_called_once_with(
+            self.context, self.session, self.instance,
+            self.host_url, ['vdi_uuid'],
+            fake_tarpipe_in, compr_level)
+        mock_stream_obj.start.assert_called_once_with()
+        fake_tarpipe_in.close.assert_called_once_with()
+
+    @mock.patch.object(vdi_handler, 'VdisToTarStream')
+    def test_start_stream_generator_abnormal_level(self, mock_stream):
+        # Verify the vdi_handler.DEFAULT_COMPRESSLEVEL should be used,
+        # if the compresslevel specified in GenerateImageStream
+        # is abnormal value (not in 1 - 9).
+        compr_level = 10
+        mock_stream_obj = mock.Mock()
+        mock_stream.return_value = mock_stream_obj
+        generator = vdi_handler.GenerateImageStream(
+            self.context, self.session, self.instance,
+            self.host_url, ['vdi_uuid'], compresslevel=compr_level)
+        fake_tarpipe_in = mock.Mock()
+
+        generator.start_image_stream_generator(fake_tarpipe_in)
+
+        mock_stream.assert_called_once_with(
+            self.context, self.session, self.instance,
+            self.host_url, ['vdi_uuid'],
+            fake_tarpipe_in, vdi_handler.DEFAULT_COMPRESSLEVEL)
+        mock_stream_obj.start.assert_called_once_with()
+        fake_tarpipe_in.close.assert_called_once_with()
+
+    @mock.patch.object(vdi_handler, 'VdisToTarStream')
+    def test_start_stream_generator_none_level(self, mock_stream):
+        # Verify the vdi_handler.DEFAULT_COMPRESSLEVEL should be used,
+        # if no compresslevel specified in GenerateImageStream.
+        mock_stream_obj = mock.Mock()
+        mock_stream.return_value = mock_stream_obj
+        generator = vdi_handler.GenerateImageStream(
+            self.context, self.session, self.instance,
+            self.host_url, ['vdi_uuid'])
+        fake_tarpipe_in = mock.Mock()
+
+        generator.start_image_stream_generator(fake_tarpipe_in)
+
+        mock_stream.assert_called_once_with(
+            self.context, self.session, self.instance,
+            self.host_url, ['vdi_uuid'],
+            fake_tarpipe_in, vdi_handler.DEFAULT_COMPRESSLEVEL)
+        mock_stream_obj.start.assert_called_once_with()
+        fake_tarpipe_in.close.assert_called_once_with()
+
 
 class VdisToTarStreamTestCase(base.TestCase):
     def setUp(self):
@@ -230,7 +293,7 @@ class VdisToTarStreamTestCase(base.TestCase):
         self.host_url = "http://fake-host.com"
         self.stream = mock.Mock()
 
-    @mock.patch.object(tarfile, 'open')
+    @mock.patch.object(tarfile.TarFile, 'gzopen')
     @mock.patch.object(tarfile, 'TarInfo')
     @mock.patch.object(vdi_handler.VdisToTarStream, '_connect_request',
                        return_value='fake-conn-resp')
@@ -254,14 +317,15 @@ class VdisToTarStreamTestCase(base.TestCase):
         vdi_uuids = ['vdi-uuid']
         vhdpipe_in = mock.Mock()
         mock_pipe.return_value = ('vhdpipe_out', vhdpipe_in)
+        compr_level = 5
         image_cmd = vdi_handler.VdisToTarStream(
             self.context, self.session, self.instance,
-            self.host_url, vdi_uuids, self.stream)
+            self.host_url, vdi_uuids, self.stream, compr_level)
 
         image_cmd.start()
 
-        mock_open.assert_called_once_with(fileobj=self.stream,
+        mock_open.assert_called_once_with(name=None, fileobj=self.stream,
-                                          mode='w|gz')
+                                          mode='w', compresslevel=compr_level)
         self.session.VDI.get_by_uuid.assert_called_once_with('vdi-uuid')
         mock_conn_req.assert_called_once_with('fake-vdi-ref')
         mock_dynDisk.get_vhd_file_size.assert_called_once_with()