Sheepdog: fix upload failure in API v2

The image_size is zero in API v2, when the client is doing a chunked
http upload. We cannot know actual image size until EOF. On-demand
resizing resolved this problem.

Closes-Bug: #1536928
Change-Id: Ia8efa4383d0ff1c8701b3ee45aeded3fca8c4339
Signed-off-by: YAMADA Hideki <yamada.hideki@lab.ntt.co.jp>

2 files changed, 56 insertions, 39 deletions

diff --git a/glance_store/_drivers/sheepdog.py b/glance_store/_drivers/sheepdog.py
index 8a40765..d24f68f 100644
--- a/glance_store/_drivers/sheepdog.py
+++ b/glance_store/_drivers/sheepdog.py
@@ -1,4 +1,5 @@
 # Copyright 2013 Taobao Inc.
+# Copyright (C) 2016 Nippon Telegraph and Telephone Corporation.
 # All Rights Reserved.
 #
 #    Licensed under the Apache License, Version 2.0 (the "License"); you may
@@ -25,6 +26,7 @@ from oslo_utils import units
 
 import glance_store
 from glance_store import capabilities
+from glance_store.common import utils
 import glance_store.driver
 from glance_store import exceptions
 from glance_store.i18n import _
@@ -76,7 +78,7 @@ class SheepdogImage(object):
 
     def get_size(self):
         """
-        Return the size of the this iamge
+        Return the size of the this image
 
         Sheepdog Usage: collie vdi list -r -a address -p port image
         """
@@ -109,6 +111,13 @@ class SheepdogImage(object):
         """
         self._run_command("create", None, str(size))
 
+    def resize(self, size):
+        """Resize this image in the Sheepdog cluster with size 'size'.
+
+        Sheepdog Usage: collie vdi create -a address -p port image size
+        """
+        self._run_command("resize", None, str(size))
+
     def delete(self):
         """
         Delete this image in the Sheepdog cluster
@@ -292,27 +301,33 @@ class Store(glance_store.driver.Store):
             'addr': self.addr,
             'port': self.port
         }, self.conf)
-        checksum = hashlib.md5()
 
         image.create(image_size)
 
         try:
-            total = left = image_size
-            while left > 0:
-                length = min(self.chunk_size, left)
-                data = image_file.read(length)
-                image.write(data, total - left, length)
-                left -= length
-                checksum.update(data)
+            offset = 0
+            checksum = hashlib.md5()
+            chunks = utils.chunkreadable(image_file, self.WRITE_CHUNKSIZE)
+            for chunk in chunks:
+                chunk_length = len(chunk)
+                # If the image size provided is zero we need to do
+                # a resize for the amount we are writing. This will
+                # be slower so setting a higher chunk size may
+                # speed things up a bit.
+                if image_size == 0:
+                    image.resize(offset + chunk_length)
+                image.write(chunk, offset, chunk_length)
+                offset += chunk_length
+                checksum.update(chunk)
                 if verifier:
-                    verifier.update(data)
+                    verifier.update(chunk)
         except Exception:
             # Note(zhiyan): clean up already received data when
             # error occurs such as ImageSizeLimitExceeded exceptions.
             with excutils.save_and_reraise_exception():
                 image.delete()
 
-        return (location.get_uri(), image_size, checksum.hexdigest(), {})
+        return (location.get_uri(), offset, checksum.hexdigest(), {})
 
     @capabilities.check
     def delete(self, location, context=None):
diff --git a/glance_store/tests/unit/test_sheepdog_store.py b/glance_store/tests/unit/test_sheepdog_store.py
index 4758349..5e6a54f 100644
--- a/glance_store/tests/unit/test_sheepdog_store.py
+++ b/glance_store/tests/unit/test_sheepdog_store.py
@@ -47,34 +47,36 @@ class TestSheepdogStore(base.StoreBaseTest,
                             'addr': 'fake_addr',
                             'port': 'fake_port'}
 
-    def test_add_image(self):
-        called_commands = []
-
-        def _fake_run_command(command, data, *params):
-            called_commands.append(command)
-
-        with mock.patch.object(sheepdog.SheepdogImage, '_run_command') as cmd:
-            cmd.side_effect = _fake_run_command
-            data = six.BytesIO(b'xx')
-            ret = self.store.add('fake_image_id', data, 2)
-            self.assertEqual(called_commands, ['list -r', 'create', 'write'])
-            self.assertEqual(ret[1], 2)
-
-    def test_cleanup_when_add_image_exception(self):
-        called_commands = []
-
-        def _fake_run_command(command, data, *params):
-            if command == 'write':
-                raise exceptions.BackendException
-            else:
-                called_commands.append(command)
-
-        with mock.patch.object(sheepdog.SheepdogImage, '_run_command') as cmd:
-            cmd.side_effect = _fake_run_command
-            data = six.BytesIO(b'xx')
-            self.assertRaises(exceptions.BackendException, self.store.add,
-                              'fake_image_id', data, 2)
-            self.assertTrue('delete' in called_commands)
+    @mock.patch.object(sheepdog.SheepdogImage, 'write')
+    @mock.patch.object(sheepdog.SheepdogImage, 'create')
+    @mock.patch.object(sheepdog.SheepdogImage, 'exist')
+    def test_add_image(self, mock_exist, mock_create, mock_write):
+        data = six.BytesIO(b'xx')
+        mock_exist.return_value = False
+
+        (uri, size, checksum, loc) = self.store.add('fake_image_id', data, 2)
+
+        mock_exist.assert_called_once_with()
+        mock_create.assert_called_once_with(2)
+        mock_write.assert_called_once_with(b'xx', 0, 2)
+
+    @mock.patch.object(sheepdog.SheepdogImage, 'delete')
+    @mock.patch.object(sheepdog.SheepdogImage, 'write')
+    @mock.patch.object(sheepdog.SheepdogImage, 'create')
+    @mock.patch.object(sheepdog.SheepdogImage, 'exist')
+    def test_cleanup_when_add_image_exception(self, mock_exist, mock_create,
+                                              mock_write, mock_delete):
+        data = six.BytesIO(b'xx')
+        mock_exist.return_value = False
+        mock_write.side_effect = exceptions.BackendException
+
+        self.assertRaises(exceptions.BackendException, self.store.add,
+                          'fake_image_id', data, 2)
+
+        mock_exist.assert_called_once_with()
+        mock_create.assert_called_once_with(2)
+        mock_write.assert_called_once_with(b'xx', 0, 2)
+        mock_delete.assert_called_once_with()
 
     def test_add_duplicate_image(self):
         def _fake_run_command(command, data, *params):