Add support for streaming raw images directly onto the disk

This patch adds support for streaming a raw image directly onto the disk,
that means no more time spent writing the image to a tmpfs partition prior
to copying it to the disk. Checksum computation is also done as the image
is being streamed. Streaming raw images is disabled by default, however
this behavior can be enabled by passing a key called "stream_raw_images"
with the value of True to the prepare_image() command of IPA.

For non-raw images this may not be possible, not sure about all image
file formats, but common types such as qcow2 requires random access to
the image file in order to be converted to raw.

Closes-Bug: #1505685
Change-Id: Iddf67907bc9b54bbd3065a97064cb5a3602cfe18

ironic_python_agent/extensions/standby.py

@@ -220,6 +220,30 @@ class StandbyExtension(base.BaseAgentExtension):
 
         self.cached_image_id = None
 
+    def _cache_and_write_image(self, image_info, device):
+        _download_image(image_info)
+        _write_image(image_info, device)
+        self.cached_image_id = image_info['id']
+
+    def _stream_raw_image_onto_device(self, image_info, device):
+        starttime = time.time()
+        image_download = ImageDownload(image_info, time_obj=starttime)
+
+        with open(device, 'wb+') as f:
+            try:
+                for chunk in image_download:
+                    f.write(chunk)
+            except Exception as e:
+                msg = 'Unable to write image to device {0}. Error: {1}'.format(
+                      device, str(e))
+                raise errors.ImageDownloadError(image_info['id'], msg)
+
+        totaltime = time.time() - starttime
+        LOG.info("Image streamed onto device {0} in {1} "
+                 "seconds".format(device, totaltime))
+        # Verify if the checksum of the streamed image is correct
+        _verify_image(image_info, device, image_download.md5sum())
+
     @base.async_command('cache_image', _validate_image_info)
     def cache_image(self, image_info=None, force=False):
         LOG.debug('Caching image %s', image_info['id'])
@@ -230,9 +254,7 @@ class StandbyExtension(base.BaseAgentExtension):
         if self.cached_image_id != image_info['id'] or force:
             LOG.debug('Already had %s cached, overwriting',
                       self.cached_image_id)
-            _download_image(image_info)
-            _write_image(image_info, device)
-            self.cached_image_id = image_info['id']
+            self._cache_and_write_image(image_info, device)
             result_msg = 'image ({0}) cached to device {1}'
 
         msg = result_msg.format(image_info['id'], device)
@@ -246,13 +268,19 @@ class StandbyExtension(base.BaseAgentExtension):
         LOG.debug('Preparing image %s', image_info['id'])
         device = hardware.dispatch_to_managers('get_os_install_device')
 
+        disk_format = image_info.get('disk_format')
+        stream_raw_images = image_info.get('stream_raw_images', False)
         # don't write image again if already cached
         if self.cached_image_id != image_info['id']:
-            LOG.debug('Already had %s cached, overwriting',
-                      self.cached_image_id)
-            _download_image(image_info)
-            _write_image(image_info, device)
-            self.cached_image_id = image_info['id']
+
+            if self.cached_image_id is not None:
+                LOG.debug('Already had %s cached, overwriting',
+                          self.cached_image_id)
+
+            if stream_raw_images and disk_format == 'raw':
+                self._stream_raw_image_onto_device(image_info, device)
+            else:
+                self._cache_and_write_image(image_info, device)
 
         if configdrive is not None:
             _write_configdrive_to_partition(configdrive, device)

ironic_python_agent/tests/unit/extensions/test_standby.py

@@ -467,6 +467,52 @@ class TestStandbyExtension(test_base.BaseTestCase):
                       ).format(image_info['id'], 'manager')
         self.assertEqual(cmd_result, async_result.command_result['result'])
 
+    @mock.patch(('ironic_python_agent.extensions.standby.'
+                 '_write_configdrive_to_partition'),
+                autospec=True)
+    @mock.patch('ironic_python_agent.hardware.dispatch_to_managers',
+                autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby.StandbyExtension'
+                '._cache_and_write_image', autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby.StandbyExtension'
+                '._stream_raw_image_onto_device', autospec=True)
+    def _test_prepare_image_raw(self, image_info, stream_mock,
+                                cache_write_mock, dispatch_mock,
+                                configdrive_copy_mock):
+        dispatch_mock.return_value = '/dev/foo'
+        configdrive_copy_mock.return_value = None
+
+        async_result = self.agent_extension.prepare_image(
+            image_info=image_info,
+            configdrive=None
+        )
+        async_result.join()
+
+        dispatch_mock.assert_called_once_with('get_os_install_device')
+        self.assertFalse(configdrive_copy_mock.called)
+
+        # Assert we've streamed the image or not
+        if image_info['stream_raw_images']:
+            stream_mock.assert_called_once_with(mock.ANY, image_info,
+                                                '/dev/foo')
+            self.assertFalse(cache_write_mock.called)
+        else:
+            cache_write_mock.assert_called_once_with(mock.ANY, image_info,
+                                                     '/dev/foo')
+            self.assertFalse(stream_mock.called)
+
+    def test_prepare_image_raw_stream_true(self):
+        image_info = _build_fake_image_info()
+        image_info['disk_format'] = 'raw'
+        image_info['stream_raw_images'] = True
+        self._test_prepare_image_raw(image_info)
+
+    def test_prepare_image_raw_and_stream_false(self):
+        image_info = _build_fake_image_info()
+        image_info['disk_format'] = 'raw'
+        image_info['stream_raw_images'] = False
+        self._test_prepare_image_raw(image_info)
+
     @mock.patch('ironic_python_agent.utils.execute', autospec=True)
     def test_run_image(self, execute_mock):
         script = standby._path_to_script('shell/shutdown.sh')
@@ -515,6 +561,62 @@ class TestStandbyExtension(test_base.BaseTestCase):
         execute_mock.assert_called_once_with(*command, check_exit_code=[0])
         self.assertEqual('FAILED', failed_result.command_status)
 
+    @mock.patch('ironic_python_agent.extensions.standby._write_image',
+                autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby._download_image',
+                autospec=True)
+    def test_cache_and_write_image(self, download_mock, write_mock):
+        image_info = _build_fake_image_info()
+        device = '/dev/foo'
+        self.agent_extension._cache_and_write_image(image_info, device)
+        download_mock.assert_called_once_with(image_info)
+        write_mock.assert_called_once_with(image_info, device)
+
+    @mock.patch('hashlib.md5')
+    @mock.patch(OPEN_FUNCTION_NAME)
+    @mock.patch('requests.get')
+    def test_stream_raw_image_onto_device(self, requests_mock, open_mock,
+                                          md5_mock):
+        image_info = _build_fake_image_info()
+        response = requests_mock.return_value
+        response.status_code = 200
+        response.iter_content.return_value = ['some', 'content']
+        file_mock = mock.Mock()
+        open_mock.return_value.__enter__.return_value = file_mock
+        file_mock.read.return_value = None
+        hexdigest_mock = md5_mock.return_value.hexdigest
+        hexdigest_mock.return_value = image_info['checksum']
+
+        self.agent_extension._stream_raw_image_onto_device(image_info,
+                                                           '/dev/foo')
+        requests_mock.assert_called_once_with(image_info['urls'][0],
+                                              stream=True, proxies={})
+        expected_calls = [mock.call('some'), mock.call('content')]
+        file_mock.write.assert_has_calls(expected_calls)
+
+    @mock.patch('hashlib.md5')
+    @mock.patch(OPEN_FUNCTION_NAME)
+    @mock.patch('requests.get')
+    def test_stream_raw_image_onto_device_write_error(self, requests_mock,
+                                                      open_mock, md5_mock):
+        image_info = _build_fake_image_info()
+        response = requests_mock.return_value
+        response.status_code = 200
+        response.iter_content.return_value = ['some', 'content']
+        file_mock = mock.Mock()
+        open_mock.return_value.__enter__.return_value = file_mock
+        file_mock.write.side_effect = Exception('Surprise!!!1!')
+        hexdigest_mock = md5_mock.return_value.hexdigest
+        hexdigest_mock.return_value = image_info['checksum']
+
+        self.assertRaises(errors.ImageDownloadError,
+                          self.agent_extension._stream_raw_image_onto_device,
+                          image_info, '/dev/foo')
+        requests_mock.assert_called_once_with(image_info['urls'][0],
+                                              stream=True, proxies={})
+        # Assert write was only called once and failed!
+        file_mock.write.assert_called_once_with('some')
+
 
 class TestImageDownload(test_base.BaseTestCase):