Add image_stage_total quota enforcement

This makes us enforce a quota on the amount of data a user has in
staging.

Partially-implements: blueprint glance-unified-quotas

Change-Id: I3cca4e589adc0aec138e5933c311aefd69ccee51

glance/api/v2/image_data.py

@@ -308,6 +308,13 @@ class ImageDataController(object):
 
     @utils.mutating
     def stage(self, req, image_id, data, size):
+        try:
+            ks_quota.enforce_image_staging_total(req.context,
+                                                 req.context.owner)
+        except exception.LimitExceeded as e:
+            raise webob.exc.HTTPRequestEntityTooLarge(explanation=str(e),
+                                                      request=req)
+
         image_repo = self.gateway.get_repo(req.context)
         image = None
 

glance/async_/flows/api_image_import.py

@@ -875,5 +875,12 @@ def get_flow(**kwargs):
                      stores, action_wrapper,
                      ks_quota.enforce_image_size_total,
                      delta=image_size)
+    elif import_method in ('copy-image', 'web-download'):
+        # The copy-image and web-download methods will use staging space to
+        # do their work, so check that quota.
+        assert_quota(kwargs['context'], task_repo, task_id,
+                     stores, action_wrapper,
+                     ks_quota.enforce_image_staging_total,
+                     delta=image_size)
 
     return flow

glance/quota/keystone.py

@@ -28,6 +28,7 @@ LOG = logging.getLogger(__name__)
 limit.opts.register_opts(CONF)
 
 QUOTA_IMAGE_SIZE_TOTAL = 'image_size_total'
+QUOTA_IMAGE_STAGING_TOTAL = 'image_stage_total'
 
 
 def _enforce_some(context, project_id, quota_value_fns, deltas):
@@ -99,3 +100,15 @@ def enforce_image_size_total(context, project_id, delta=0):
         context, project_id, QUOTA_IMAGE_SIZE_TOTAL,
         lambda: db.user_get_storage_usage(context, project_id) // units.Mi,
         delta=delta)
+
+
+def enforce_image_staging_total(context, project_id, delta=0):
+    """Enforce the image_stage_total quota.
+
+    This enforces the total size of all images stored in staging areas
+    for the supplied project_id.
+    """
+    _enforce_one(
+        context, project_id, QUOTA_IMAGE_STAGING_TOTAL,
+        lambda: db.user_get_staging_usage(context, project_id) // units.Mi,
+        delta=delta)

glance/tests/functional/v2/test_images.py

@@ -7140,8 +7140,9 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
         self.assertEqual('success', task['status'])
 
     def test_copy(self):
-        # Set a quota of 5MiB
-        self.set_limit({'image_size_total': 5})
+        # Set a size quota of 5MiB, with more staging quota than we need.
+        self.set_limit({'image_size_total': 5,
+                        'image_stage_total': 15})
         self.start_server()
 
         # First import of 3MiB is good
@@ -7154,7 +7155,61 @@ class TestKeystoneQuotas(functional.SynchronousAPIBase):
         req = self._import_copy(image_id, ['store2'])
         self.assertEqual(202, req.status_code)
         self._wait_for_import(image_id)
+        self.assertEqual('success', self._get_latest_task(image_id)['status'])
 
-        # Third copy should fail because we're over quota
+        # Third copy should fail because we're over total size quota.
         req = self._import_copy(image_id, ['store3'])
         self.assertEqual(413, req.status_code)
+
+        # Set our size quota to have enough space, but restrict our
+        # staging quota to below the required size to stage the image
+        # before copy. This request should succeed, but the copy task
+        # should fail the staging quota check.
+        self.set_limit({'image_size_total': 15,
+                        'image_stage_total': 5})
+        req = self._import_copy(image_id, ['store3'])
+        self.assertEqual(202, req.status_code)
+        self._wait_for_import(image_id)
+        self.assertEqual('failure', self._get_latest_task(image_id)['status'])
+
+        # If we increase our stage quota, we should now be able to copy.
+        self.set_limit({'image_size_total': 15,
+                        'image_stage_total': 10})
+        req = self._import_copy(image_id, ['store3'])
+        self.assertEqual(202, req.status_code)
+        self._wait_for_import(image_id)
+        self.assertEqual('success', self._get_latest_task(image_id)['status'])
+
+    def test_stage(self):
+        # Set a quota of 5MiB
+        self.set_limit({'image_size_total': 15,
+                        'image_stage_total': 5})
+        self.start_server()
+
+        # Stage 6MiB, which is allowed to complete, but leaves us over
+        # quota
+        image_id = self._create_and_stage(
+            data_iter=test_utils.FakeData(6 * units.Mi))
+
+        # Second stage fails because we are out of quota
+        self._create_and_stage(expected_code=413)
+
+        # Make sure that a web-download fails to actually run.
+        image_id2 = self._create().json['id']
+        req = self._import_web_download(image_id2, ['store1'],
+                                        'http://example.com/foo.img')
+        self.assertEqual(202, req.status_code)
+        self._wait_for_import(image_id2)
+        task = self._get_latest_task(image_id2)
+        self.assertEqual('failure', task['status'])
+        self.assertIn('image_stage_total is over limit', task['message'])
+
+        # Finish importing one of the images, which should put us under quota
+        # for staging
+        req = self._import_direct(image_id, ['store1'])
+        self.assertEqual(202, req.status_code)
+        self._wait_for_import(image_id)
+
+        # Stage should now succeed because we have freed up quota
+        self._create_and_stage(
+            data_iter=test_utils.FakeData(6 * units.Mi))