Quota usage support in Cyborg

1. Introduce quote usage related tables.
2. Add reserve() and commit() function to update quota usage in DB.
3. Invoke reserve() and commit() when users create or delete acclerators.

Change-Id: I828bc6d35d08116a2b3c74baeda8876121541f8c

cyborg/api/controllers/v1/accelerators.py

@@ -230,9 +230,8 @@ class AcceleratorsController(AcceleratorsControllerBase):
     @expose.expose(None, types.uuid, status_code=http_client.NO_CONTENT)
     def delete(self, uuid):
         """Delete an accelerator.
-
         :param uuid: UUID of an accelerator.
         """
-        obj_acc = self._resource or self._get_resource(uuid)
         context = pecan.request.context
+        obj_acc = self._resource or self._get_resource(uuid)
         pecan.request.conductor_api.accelerator_delete(context, obj_acc)

cyborg/api/controllers/v1/deployables.py

@@ -27,6 +27,7 @@ from cyborg.api import expose
 from cyborg.common import exception
 from cyborg.common import policy
 from cyborg import objects
+from cyborg.quota import QUOTAS
 
 
 class Deployable(base.APIBase):
@@ -226,13 +227,29 @@ class DeployablesController(base.CyborgController):
         :param patch: a json PATCH document to apply to this deployable.
         """
         context = pecan.request.context
+        reservations = None
+
         obj_dep = objects.Deployable.get(context, uuid)
         try:
+            # TODO(xinran): need more discussion on quota's granularity.
+            # Now we count by board.
+            for p in patch:
+                if p["path"] == "/instance_uuid" and p["op"] == "replace":
+                    if not p["value"]:
+                        obj_dep["assignable"] = True
+                        reserve_opts = {obj_dep["board"]: -1}
+                    else:
+                        obj_dep["assignable"] = False
+                        reserve_opts = {obj_dep["board"]: 1}
+                    reservations = QUOTAS.reserve(context, reserve_opts)
             api_dep = Deployable(
                 **api_utils.apply_jsonpatch(obj_dep.as_dict(), patch))
         except api_utils.JSONPATCH_EXCEPTIONS as e:
+            QUOTAS.rollback(context, reservations, project_id=None)
             raise exception.PatchError(patch=patch, reason=e)
 
+        QUOTAS.commit(context, reservations)
+
         # Update only the fields that have changed
         for field in objects.Deployable.fields:
             try:

cyborg/common/exception.py

@@ -289,3 +289,19 @@ class InventoryInUse(InvalidInventory):
     # cyborg.services.client.report._RE_INV_IN_USE regex.
     msg_fmt = _("Inventory for '%(resource_classes)s' on "
                 "resource provider '%(resource_provider)s' in use.")
+
+
+class QuotaNotFound(NotFound):
+    message = _("Quota could not be found")
+
+
+class QuotaUsageNotFound(QuotaNotFound):
+    message = _("Quota usage for project %(project_id)s could not be found.")
+
+
+class QuotaResourceUnknown(QuotaNotFound):
+    message = _("Unknown quota resources %(unknown)s.")
+
+
+class InvalidReservationExpiration(Invalid):
+    message = _("Invalid reservation expiration %(expire)s.")

cyborg/db/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2010 United States Government as represented by the
+# Administrator of the National Aeronautics and Space Administration.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+"""
+DB abstraction for Cyborg
+"""
+
+from cyborg.db.api import *  # noqa

cyborg/db/api.py

@@ -123,3 +123,13 @@ class Connection(object):
     @abc.abstractmethod
     def attribute_delete(self, context, uuid):
         """Delete an attribute."""
+
+    @abc.abstractmethod
+    def quota_reserve(self, context, resources, deltas, expire,
+                      until_refresh, max_age, project_id=None,
+                      is_allocated_reserve=False):
+        """Check quotas and create appropriate reservations."""
+
+    @abc.abstractmethod
+    def reservation_commit(self, context, reservations, project_id=None):
+        """Check quotas and create appropriate reservations."""

cyborg/db/sqlalchemy/alembic/versions/d6f033d8fa5b_add_quota_related_tables.py

@@ -0,0 +1,70 @@
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""add-quota-related-tables
+
+Revision ID: d6f033d8fa5b
+Revises: f50980397351
+Create Date: 2018-04-28 03:07:06.857245
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = 'd6f033d8fa5b'
+down_revision = 'f50980397351'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        'quota_usages',
+        sa.Column('created_at', sa.DateTime(), nullable=True),
+        sa.Column('updated_at', sa.DateTime(), nullable=True),
+        sa.Column('id', sa.Integer(), nullable=False),
+        sa.Column('project_id', sa.String(length=255), nullable=True),
+        sa.Column('user_id', sa.String(length=255), nullable=True),
+        sa.Column('resource', sa.String(length=255), nullable=False),
+        sa.Column('in_use', sa.Integer(), nullable=False),
+        sa.Column('reserved', sa.Integer(), nullable=False),
+        sa.Column('until_refresh', sa.Integer(), nullable=True),
+        sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index('ix_quota_usages_project_id', 'quota_usages',
+                    ['project_id'], unique=False)
+    op.create_index('ix_quota_usages_user_id', 'quota_usages', ['user_id'],
+                    unique=False)
+
+    op.create_table(
+        'reservations',
+        sa.Column('created_at', sa.DateTime(), nullable=True),
+        sa.Column('updated_at', sa.DateTime(), nullable=True),
+        sa.Column('id', sa.Integer(), nullable=False),
+        sa.Column('uuid', sa.String(length=36), nullable=False),
+        sa.Column('usage_id', sa.Integer(), nullable=False),
+        sa.Column('project_id', sa.String(length=255), nullable=True),
+        sa.Column('user_id', sa.String(length=255), nullable=True),
+        sa.Column('resource', sa.String(length=255), nullable=True),
+        sa.Column('delta', sa.Integer(), nullable=False),
+        sa.Column('expire', sa.DateTime(), nullable=True),
+        sa.ForeignKeyConstraint(['usage_id'], ['quota_usages.id'], ),
+        sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index('ix_reservations_project_id', 'reservations',
+                    ['project_id'], unique=False)
+    op.create_index('ix_reservations_user_id', 'reservations', ['user_id'],
+                    unique=False)
+    op.create_index('reservations_uuid_idx', 'reservations', ['uuid'],
+                    unique=False)
+    # ### end Alembic commands ###

cyborg/db/sqlalchemy/api.py

@@ -17,6 +17,7 @@
 
 import threading
 import copy
+import uuid
 
 from oslo_db import api as oslo_db_api
 from oslo_db import exception as db_exc
@@ -24,8 +25,12 @@ from oslo_db.sqlalchemy import enginefacade
 from oslo_db.sqlalchemy import utils as sqlalchemyutils
 from oslo_log import log
 from oslo_utils import strutils
+from oslo_utils import timeutils
 from oslo_utils import uuidutils
+from sqlalchemy.orm import load_only
 from sqlalchemy.orm.exc import NoResultFound
+from sqlalchemy.sql import func
+
 
 from cyborg.common import exception
 from cyborg.common.i18n import _
@@ -37,6 +42,8 @@ from sqlalchemy import and_
 _CONTEXT = threading.local()
 LOG = log.getLogger(__name__)
 
+main_context_manager = enginefacade.transaction_context()
+
 
 def get_backend():
     """The backend is this module itself."""
@@ -51,6 +58,11 @@ def _session_for_write():
     return enginefacade.writer.using(_CONTEXT)
 
 
+def get_session(use_slave=False, **kwargs):
+    return main_context_manager._factory.get_legacy_facade().get_session(
+        use_slave=use_slave, **kwargs)
+
+
 def model_query(context, model, *args, **kwargs):
     """Query helper for simpler session usage.
 
@@ -498,6 +510,185 @@ class Connection(api.Connection):
             if count != 1:
                 raise exception.AttributeNotFound(uuid=uuid)
 
+    def _get_quota_usages(self, context, project_id, resources=None):
+        # Broken out for testability
+        query = model_query(context, models.QuotaUsage,).filter_by(
+            project_id=project_id)
+        if resources:
+            query = query.filter(models.QuotaUsage.resource.in_(
+                list(resources)))
+        rows = query.order_by(models.QuotaUsage.id.asc()). \
+            with_for_update().all()
+        return {row.resource: row for row in rows}
+
+    def _quota_usage_create(self, project_id, resource, until_refresh,
+                            in_use, reserved, session=None):
+
+        quota_usage_ref = models.QuotaUsage()
+        quota_usage_ref.project_id = project_id
+        quota_usage_ref.resource = resource
+        quota_usage_ref.in_use = in_use
+        quota_usage_ref.reserved = reserved
+        quota_usage_ref.until_refresh = until_refresh
+        quota_usage_ref.save(session=session)
+
+        return quota_usage_ref
+
+    def _reservation_create(self, uuid, usage, project_id, resource, delta,
+                            expire, session=None):
+        usage_id = usage['id'] if usage else None
+        reservation_ref = models.Reservation()
+        reservation_ref.uuid = uuid
+        reservation_ref.usage_id = usage_id
+        reservation_ref.project_id = project_id
+        reservation_ref.resource = resource
+        reservation_ref.delta = delta
+        reservation_ref.expire = expire
+        reservation_ref.save(session=session)
+        return reservation_ref
+
+    def _get_reservation_resources(self, context, reservation_ids):
+        """Return the relevant resources by reservations."""
+
+        reservations = model_query(context, models.Reservation). \
+            options(load_only('resource')). \
+            filter(models.Reservation.uuid.in_(reservation_ids)). \
+            all()
+        return {r.resource for r in reservations}
+
+    def _quota_reservations(self, session, context, reservations):
+        """Return the relevant reservations."""
+
+        # Get the listed reservations
+        return model_query(context, models.Reservation). \
+            filter(models.Reservation.uuid.in_(reservations)). \
+            with_lockmode('update'). \
+            all()
+
+    def quota_reserve(self, context, resources, deltas, expire,
+                      until_refresh, max_age, project_id=None,
+                      is_allocated_reserve=False):
+        """ Create reservation record in DB according to params"""
+        with _session_for_write() as session:
+            if project_id is None:
+                project_id = context.project_id
+            usages = self._get_quota_usages(context, project_id,
+                                            resources=deltas.keys())
+            work = set(deltas.keys())
+            while work:
+                resource = work.pop()
+
+                # Do we need to refresh the usage?
+                refresh = False
+                # create quota usage in DB if there is no record of this type
+                # of resource
+                if resource not in usages:
+                    usages[resource] = self._quota_usage_create(project_id,
+                                                                resource,
+                                                                until_refresh
+                                                                or None,
+                                                                in_use=0,
+                                                                reserved=0,
+                                                                session=session
+                                                                )
+                    refresh = True
+                elif usages[resource].in_use < 0:
+                    # Negative in_use count indicates a desync, so try to
+                    # heal from that...
+                    refresh = True
+                elif usages[resource].until_refresh is not None:
+                    usages[resource].until_refresh -= 1
+                    if usages[resource].until_refresh <= 0:
+                        refresh = True
+                elif max_age and usages[resource].updated_at is not None and (
+                    (timeutils.utcnow() -
+                        usages[resource].updated_at).total_seconds() >=
+                        max_age):
+                    refresh = True
+
+                # refresh the usage
+                if refresh:
+                    # Grab the sync routine
+                    updates = self._sync_acc_res(context, resource, project_id)
+                    for res, in_use in updates.items():
+                        # Make sure we have a destination for the usage!
+                        if res not in usages:
+                            usages[res] = self._quota_usage_create(
+                                project_id,
+                                res,
+                                until_refresh or None,
+                                in_use=0,
+                                reserved=0,
+                                session=session
+                            )
+
+                        # Update the usage
+                        usages[res].in_use = in_use
+                        usages[res].until_refresh = until_refresh or None
+
+                        # Because more than one resource may be refreshed
+                        # by the call to the sync routine, and we don't
+                        # want to double-sync, we make sure all refreshed
+                        # resources are dropped from the work set.
+                        work.discard(res)
+
+                        # NOTE(Vek): We make the assumption that the sync
+                        #            routine actually refreshes the
+                        #            resources that it is the sync routine
+                        #            for.  We don't check, because this is
+                        #            a best-effort mechanism.
+
+            unders = [r for r, delta in deltas.items()
+                      if delta < 0 and delta + usages[r].in_use < 0]
+            reservations = []
+            for resource, delta in deltas.items():
+                usage = usages[resource]
+                reservation = self._reservation_create(
+                    str(uuid.uuid4()), usage, project_id, resource,
+                    delta, expire, session=session)
+                reservations.append(reservation.uuid)
+                usages[resource].reserved += delta
+            session.flush()
+        if unders:
+            LOG.warning("Change will make usage less than 0 for the following "
+                        "resources: %s", unders)
+        return reservations
+
+    def _sync_acc_res(self, context, resource, project_id):
+        """Quota sync funciton"""
+        res_in_use = self._accelerator_data_get_for_project(context, resource,
+                                                            project_id)
+        return {resource: res_in_use}
+
+    def _accelerator_data_get_for_project(self, context, resource, project_id):
+        """Return the number of resource which is being used by a project"""
+        query = model_query(context, models.Accelerator).\
+            filter_by(project_id=project_id).filter_by(device_type=resource)
+
+        return query.count()
+
+    def _dict_with_usage_id(self, usages):
+        return {row.id: row for row in usages.values()}
+
+    def reservation_commit(self, context, reservations, project_id=None):
+        """Commit quota reservation to quota usage table"""
+        with _session_for_write() as session:
+            quota_usage = self._get_quota_usages(
+                context, project_id,
+                resources=self._get_reservation_resources(context,
+                                                          reservations))
+            usages = self._dict_with_usage_id(quota_usage)
+
+            for reservation in self._quota_reservations(session, context,
+                                                        reservations):
+
+                usage = usages[reservation.usage_id]
+                if reservation.delta >= 0:
+                    usage.reserved -= reservation.delta
+                usage.in_use += reservation.delta
+                session.flush()
+                reservation.delete(session=session)
+
     def process_sort_params(self, sort_keys, sort_dirs,
                             default_keys=['created_at', 'id'],
                             default_dir='asc'):

cyborg/db/sqlalchemy/models.py

@@ -17,11 +17,14 @@
 
 from oslo_db import options as db_options
 from oslo_db.sqlalchemy import models
+from oslo_utils import timeutils
 import six.moves.urllib.parse as urlparse
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import Column, String, Integer, Boolean, ForeignKey, Index
 from sqlalchemy import Text
 from sqlalchemy import schema
+from sqlalchemy import DateTime
+from sqlalchemy import orm
 
 from cyborg.common import paths
 from cyborg.conf import CONF
@@ -48,6 +51,18 @@ class CyborgBase(models.TimestampMixin, models.ModelBase):
             d[c.name] = self[c.name]
         return d
 
+    @staticmethod
+    def delete_values():
+        return {'deleted': True,
+                'deleted_at': timeutils.utcnow()}
+
+    def delete(self, session):
+        """Delete this object."""
+        updated_values = self.delete_values()
+        self.update(updated_values)
+        self.save(session=session)
+        return updated_values
+
 
 Base = declarative_base(cls=CyborgBase)
 
@@ -124,3 +139,54 @@ class Attribute(Base):
                            nullable=False)
     key = Column(Text, nullable=False)
     value = Column(Text, nullable=False)
+
+
+class QuotaUsage(Base):
+    """Represents the current usage for a given resource."""
+
+    __tablename__ = 'quota_usages'
+    __table_args__ = (
+        Index('ix_quota_usages_project_id', 'project_id'),
+        Index('ix_quota_usages_user_id', 'user_id'),
+    )
+    id = Column(Integer, primary_key=True)
+
+    project_id = Column(String(255))
+    user_id = Column(String(255))
+    resource = Column(String(255), nullable=False)
+
+    in_use = Column(Integer, nullable=False)
+    reserved = Column(Integer, nullable=False)
+
+    @property
+    def total(self):
+        return self.in_use + self.reserved
+
+    until_refresh = Column(Integer)
+
+
+class Reservation(Base):
+    """Represents a resource reservation for quotas."""
+
+    __tablename__ = 'reservations'
+    __table_args__ = (
+        Index('ix_reservations_project_id', 'project_id'),
+        Index('reservations_uuid_idx', 'uuid'),
+        Index('ix_reservations_user_id', 'user_id'),
+    )
+    id = Column(Integer, primary_key=True, nullable=False)
+    uuid = Column(String(36), nullable=False)
+
+    usage_id = Column(Integer, ForeignKey('quota_usages.id'), nullable=False)
+
+    project_id = Column(String(255))
+    user_id = Column(String(255))
+    resource = Column(String(255))
+
+    delta = Column(Integer, nullable=False)
+    expire = Column(DateTime)
+
+    usage = orm.relationship(
+        "QuotaUsage",
+        foreign_keys=usage_id,
+        primaryjoin=usage_id == QuotaUsage.id)

cyborg/quota.py

@@ -0,0 +1,190 @@
+# Copyright 2018 Intel, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from oslo_config import cfg
+from oslo_log import log as logging
+from oslo_utils import timeutils
+import six
+
+from cyborg import db as db_api
+from cyborg.common import exception
+
+LOG = logging.getLogger(__name__)
+
+quota_opts = [
+    cfg.IntOpt('reservation_expire',
+               default=86400,
+               help='Number of seconds until a reservation expires'),
+    cfg.IntOpt('until_refresh',
+               default=0,
+               help='Count of reservations until usage is refreshed'),
+    cfg.StrOpt('quota_driver',
+               default="cyborg.quota.DbQuotaDriver",
+               help='Default driver to use for quota checks'),
+    cfg.IntOpt('quota_fpgas',
+               default=10,
+               help='Total amount of fpga allowed per project'),
+    cfg.IntOpt('quota_gpus',
+               default=10,
+               help='Total amount of storage allowed per project'),
+    cfg.IntOpt('max_age',
+               default=0,
+               help='Number of seconds between subsequent usage refreshes')
+    ]
+
+CONF = cfg.CONF
+CONF.register_opts(quota_opts)
+
+
+class QuotaEngine(object):
+    """Represent the set of recognized quotas."""
+
+    def __init__(self, quota_driver_class=None):
+        """Initialize a Quota object."""
+
+        self._resources = {}
+        self._driver = DbQuotaDriver()
+
+    def register_resource(self, resource):
+        """Register a resource."""
+        self._resources[resource.name] = resource
+
+    def register_resources(self, resources):
+        """Register a list of resources."""
+        for resource in resources:
+            self.register_resource(resource)
+
+    def reserve(self, context, deltas, expire=None, project_id=None):
+        """Check quotas and reserve resources.
+
+        For counting quotas--those quotas for which there is a usage
+        synchronization function--this method checks quotas against
+        current usage and the desired deltas.  The deltas are given as
+        keyword arguments, and current usage and other reservations
+        are factored into the quota check.
+
+        This method will raise a QuotaResourceUnknown exception if a
+        given resource is unknown or if it does not have a usage
+        synchronization function.
+
+        If any of the proposed values is over the defined quota, an
+        OverQuota exception will be raised with the sorted list of the
+        resources which are too high.  Otherwise, the method returns a
+        list of reservation UUIDs which were created.
+
+        :param context: The request context, for access checks.
+        :param expire: An optional parameter specifying an expiration
+                       time for the reservations.  If it is a simple
+                       number, it is interpreted as a number of
+                       seconds and added to the current time; if it is
+                       a datetime.timedelta object, it will also be
+                       added to the current time.  A datetime.datetime
+                       object will be interpreted as the absolute
+                       expiration time.  If None is specified, the
+                       default expiration time set by
+                       --default-reservation-expire will be used (this
+                       value will be treated as a number of seconds).
+        :param project_id: Specify the project_id if current context
+                           is admin and admin wants to impact on
+                           common user's tenant.
+        """
+        if not project_id:
+            project_id = context.project_id
+        reservations = self._driver.reserve(context, self._resources, deltas,
+                                            expire=expire,
+                                            project_id=project_id)
+
+        LOG.debug("Created reservations %s", reservations)
+
+        return reservations
+
+    def commit(self, context, reservations, project_id=None):
+        """Commit reservations.
+
+        :param context: The request context, for access checks.
+        :param reservations: A list of the reservation UUIDs, as
+                             returned by the reserve() method.
+        :param project_id: Specify the project_id if current context
+                           is admin and admin wants to impact on
+                           common user's tenant.
+        """
+        project_id = context.project_id
+        try:
+            self._driver.commit(context, reservations, project_id=project_id)
+        except Exception:
+            # NOTE(Vek): Ignoring exceptions here is safe, because the
+            # usage resynchronization and the reservation expiration
+            # mechanisms will resolve the issue.  The exception is
+            # logged, however, because this is less than optimal.
+            LOG.exception("Failed to commit reservations %s", reservations)
+
+    def rollback(self, context, reservations, project_id=None):
+        pass
+
+
+class DbQuotaDriver(object):
+    """Driver to perform check to enforcement of quotas.
+
+    Also allows to obtain quota information.
+    The default driver utilizes the local database.
+    """
+    dbapi = db_api.get_instance()
+
+    def reserve(self, context, resources, deltas, expire=None,
+                project_id=None):
+        # Set up the reservation expiration
+        if expire is None:
+            expire = CONF.reservation_expire
+        if isinstance(expire, six.integer_types):
+            expire = datetime.timedelta(seconds=expire)
+        if isinstance(expire, datetime.timedelta):
+            expire = timeutils.utcnow() + expire
+        if not isinstance(expire, datetime.datetime):
+            raise exception.InvalidReservationExpiration(expire=expire)
+
+        # If project_id is None, then we use the project_id in context
+        if project_id is None:
+            project_id = context.project_id
+
+        return self._reserve(context, resources, deltas, expire,
+                             project_id)
+
+    def _reserve(self, context, resources, deltas, expire, project_id):
+        return self.dbapi.quota_reserve(context, resources, deltas, expire,
+                                        CONF.until_refresh, CONF.max_age,
+                                        project_id=project_id)
+
+    def commit(self, context, reservations, project_id=None):
+        """Commit reservations.
+
+        :param context: The request context, for access checks.
+        :param reservations: A list of the reservation UUIDs, as
+                             returned by the reserve() method.
+        :param project_id: Specify the project_id if current context
+                           is admin and admin wants to impact on
+                           common user's tenant.
+        """
+
+        try:
+            self.dbapi.reservation_commit(context, reservations,
+                                          project_id=project_id)
+        except Exception:
+            # NOTE(Vek): Ignoring exceptions here is safe, because the
+            # usage resynchronization and the reservation expiration
+            # mechanisms will resolve the issue.  The exception is
+            # logged, however, because this is less than optimal.
+            LOG.exception("Failed to commit reservations %s", reservations)
+
+QUOTAS = QuotaEngine()

cyborg/tests/unit/db/test_db_api.py

@@ -0,0 +1,129 @@
+# Copyright 2018 Intel, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+"""Unit tests for the DB api."""
+
+import datetime
+from cyborg.tests.unit.db import base
+from cyborg.db import api as dbapi
+from cyborg.db.sqlalchemy import api as sqlalchemyapi
+
+
+def _quota_reserve(context, project_id):
+    """Create sample QuotaUsage and Reservation objects.
+
+    There is no method db.quota_usage_create(), so we have to use
+    db.quota_reserve() for creating QuotaUsage objects.
+
+    Returns reservations uuids.
+
+    """
+    sqlalchemy_api = sqlalchemyapi.get_backend()
+    resources = {}
+    deltas = {}
+    for i, resource in enumerate(('fpga', 'gpu')):
+        deltas[resource] = i + 1
+    return sqlalchemy_api.quota_reserve(
+        context, resources, deltas,
+        datetime.datetime.utcnow(), datetime.datetime.utcnow(),
+        datetime.timedelta(days=1), project_id
+    )
+
+
+class DBAPIQuotaUsageTestCase(base.DbTestCase):
+
+    """Tests for db.api.quota_usage_* methods."""
+
+    def test_quota_reserve(self):
+        sqlalchemy_api = sqlalchemyapi.get_backend()
+        reservations = _quota_reserve(self.context, 'project1')
+        self.assertEqual(2, len(reservations))
+        quota_usages = sqlalchemy_api._get_quota_usages(self.context,
+                                                        'project1')
+        result = {'project_id': "project1"}
+        for k, v in quota_usages.items():
+            result[v.resource] = dict(in_use=v.in_use, reserved=v.reserved)
+
+        self.assertEqual({'project_id': 'project1',
+                          'gpu': {'reserved': 2, 'in_use': 0},
+                          'fpga': {'reserved': 1, 'in_use': 0}},
+                         result)
+
+    def test__get_quota_usages(self):
+        _quota_reserve(self.context, 'project1')
+        sqlalchemy_api = sqlalchemyapi.get_backend()
+        quota_usages = sqlalchemy_api._get_quota_usages(self.context,
+                                                        'project1')
+
+        self.assertEqual(['fpga', 'gpu'],
+                         sorted(quota_usages.keys()))
+
+    def test__get_quota_usages_with_resources(self):
+        _quota_reserve(self.context, 'project1')
+        sqlalchemy_api = sqlalchemyapi.get_backend()
+        quota_usage = sqlalchemy_api._get_quota_usages(
+            self.context, 'project1', resources=['gpu'])
+
+        self.assertEqual(['gpu'], list(quota_usage.keys()))
+
+
+class DBAPIReservationTestCase(base.DbTestCase):
+
+    """Tests for db.api.reservation_* methods."""
+
+    def setUp(self):
+        super(DBAPIReservationTestCase, self).setUp()
+        self.values = {
+            'uuid': 'sample-uuid',
+            'project_id': 'project1',
+            'resource': 'resource',
+            'delta': 42,
+            'expire': (datetime.datetime.utcnow() +
+                       datetime.timedelta(days=1)),
+            'usage': {'id': 1}
+        }
+
+    def test__get_reservation_resources(self):
+        sqlalchemy_api = sqlalchemyapi.get_backend()
+        reservations = _quota_reserve(self.context, 'project1')
+        expected = ['fpga', 'gpu']
+        resources = sqlalchemy_api._get_reservation_resources(
+            self.context, reservations)
+        self.assertEqual(expected, sorted(resources))
+
+    def test_reservation_commit(self):
+        db_api = dbapi.get_instance()
+        reservations = _quota_reserve(self.context, 'project1')
+        expected = {'project_id': 'project1',
+                    'fpga': {'reserved': 1, 'in_use': 0},
+                    'gpu': {'reserved': 2, 'in_use': 0},
+                    }
+        quota_usages = db_api._get_quota_usages(self.context, 'project1')
+        result = {'project_id': "project1"}
+        for k, v in quota_usages.items():
+            result[v.resource] = dict(in_use=v.in_use, reserved=v.reserved)
+
+        self.assertEqual(expected, result)
+
+        db_api.reservation_commit(self.context, reservations, 'project1')
+        expected = {'project_id': 'project1',
+                    'fpga': {'reserved': 0, 'in_use': 1},
+                    'gpu': {'reserved': 0, 'in_use': 2},
+                    }
+        quota_usages = db_api._get_quota_usages(self.context, 'project1')
+        result = {'project_id': "project1"}
+        for k, v in quota_usages.items():
+            result[v.resource] = dict(in_use=v.in_use,
+                                      reserved=v.reserved)
+        self.assertEqual(expected, result)