diff --git a/tripleo_common/image/image_export.py b/tripleo_common/image/image_export.py index b67a7103a..d289f87c8 100644 --- a/tripleo_common/image/image_export.py +++ b/tripleo_common/image/image_export.py @@ -88,6 +88,12 @@ def export_stream(target_url, layer, layer_stream, verify_digest=True): length = 0 calc_digest = hashlib.sha256() + def remove_layer(image, blob_path): + if os.path.isfile(blob_path): + os.remove(blob_path) + LOG.error('[%s] Broken layer found and removed %s' % + (image, blob_path)) + try: with open(blob_path, 'wb') as f: count = 0 @@ -102,13 +108,15 @@ def export_stream(target_url, layer, layer_stream, verify_digest=True): length += len(chunk) LOG.debug('[%s] Written %i bytes for %s' % (image, length, digest)) + except MemoryError as e: + memory_error = '[{}] Memory Error: {}'.format(image, str(e)) + LOG.error(memory_error) + remove_layer(image, blob_path) + raise MemoryError(memory_error) except Exception as e: write_error = '[{}] Write Failure: {}'.format(image, str(e)) LOG.error(write_error) - if os.path.isfile(blob_path): - os.remove(blob_path) - LOG.error('[%s] Broken layer found and removed %s' % - (image, blob_path)) + remove_layer(image, blob_path) raise IOError(write_error) else: LOG.info('[%s] Layer written successfully %s' % (image, blob_path)) diff --git a/tripleo_common/tests/image/test_image_export.py b/tripleo_common/tests/image/test_image_export.py index 1be8aa8cb..0e7169c46 100644 --- a/tripleo_common/tests/image/test_image_export.py +++ b/tripleo_common/tests/image/test_image_export.py @@ -16,6 +16,7 @@ import hashlib import io import json +import mock import os import requests import shutil @@ -105,6 +106,23 @@ class TestImageExport(base.TestCase): with open(blob_path, 'rb') as f: self.assertEqual(blob_compressed, f.read()) + @mock.patch('tripleo_common.image.image_export.open', + side_effect=MemoryError()) + def test_export_stream_memory_error(self, mock_open): + blob_data = six.b('The Blob') + blob_compressed = zlib.compress(blob_data) + calc_digest = hashlib.sha256() + calc_digest.update(blob_compressed) + + target_url = urlparse('docker://localhost:8787/t/nova-api:latest') + layer = { + 'digest': 'sha256:somethingelse' + } + calc_digest = hashlib.sha256() + layer_stream = io.BytesIO(blob_compressed) + self.assertRaises(MemoryError, image_export.export_stream, + target_url, layer, layer_stream, verify_digest=False) + def test_export_stream_verify_failed(self): blob_data = six.b('The Blob') blob_compressed = zlib.compress(blob_data)