add quota

Author: slashburygin

AGENTS.md

@@ -89,10 +89,10 @@ exordos_core/
 tox -e py310,py312,py314
 
 # Run unit tests only
-tox -e py310
+tox -e py314
 
 # Run functional tests
-tox -e py310-functional
+tox -e py314-functional
 
 # Run linters
 tox -e ruff-check    # Check code style

exordos_core/compute/dm/models.py

@@ -36,6 +36,7 @@
 from exordos_core.common import utils
 from exordos_core.common.dm import models as cm
 from exordos_core.compute import constants as nc
+from exordos_core.quota.dm.models import QuotaModelMixin
 
 if tp.TYPE_CHECKING:
     from exordos_core.compute.pool.drivers.base import AbstractPoolDriver
@@ -261,6 +262,7 @@ class UnscheduledVolume(models.ModelWithUUID, orm.SQLStorableMixin):
 class NodeSet(
     infra_models.NodeSet,
     ua_models.InstanceWithDerivativesMixin,
+    QuotaModelMixin,
     orm.SQLStorableMixin,
 ):
     __tablename__ = "compute_sets"
@@ -300,6 +302,7 @@ def set_active(self):
 
 class Node(
     infra_models.Node,
+    QuotaModelMixin,
     orm.SQLStorableWithJSONFieldsMixin,
 ):
     __tablename__ = "nodes"

exordos_core/quota/__init__.py

@@ -0,0 +1,15 @@
+#    Copyright 2026 Genesis Corporation.
+#
+#    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.

exordos_core/quota/dm/__init__.py

@@ -0,0 +1,319 @@
+#    Copyright 2026 Genesis Corporation.
+#
+#    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.
+
+#
+# Quota flow
+# ----------
+#
+# Two tables control quotas:
+#   quota_limits       — admin-defined cap per resource (e.g. net_lb → 5)
+#   quotas_reservations — per-entity reservation rows tying a quota slot to an entity UUID
+#
+# QuotaModelMixin overrides insert() and delete() on tracked entities (LB, NodeSet, Node).
+#
+# On insert():
+#   1. _quota_reserve(s) — in the caller's transaction:
+#      a. Look up the limit for this resource (self.__tablename__).
+#         If no limit is set, quota is unlimited — skip.
+#      b. Acquire a pg_advisory_xact_lock keyed by {project_id}:{resource_name}
+#         to serialize concurrent reservations for the same project+resource.
+#      c. Count existing reservations for this project+resource.
+#         If at or above the limit, raise QuotaExceededError (rolls back the transaction).
+#      d. Insert a QuotaReservation row with uuid=self.uuid — the entity's own UUID
+#         links the reservation to the entity.
+#   2. super().insert(s) — persist the entity itself in the same transaction.
+#   3. If step 2 raises, the transaction rolls back and the reservation is
+#      automatically cleaned up — no explicit release needed.
+#
+# On delete():
+#   1. _quota_release_all(s) — delete the reservation row (frees the quota slot).
+#   2. super().delete(s) — delete the entity itself.
+#
+# reconcile_quota_reservations(session):
+#   Find all reservations whose uuid no longer exists in the entity table.
+#   This handles edge cases where reservation cleanup was missed.
+#   Designed to be run periodically or on-demand as a repair mechanism.
+#
+
+import hashlib
+import logging
+import typing as tp
+import uuid as sys_uuid
+
+from restalchemy.common import exceptions as ra_e
+from restalchemy.dm import filters as dm_filters
+from restalchemy.dm import models
+from restalchemy.dm import properties
+from restalchemy.dm import types
+from restalchemy.storage import exceptions as rs_e
+from restalchemy.storage.sql import orm
+
+LOG = logging.getLogger(__name__)
+
+
+class QuotaExceededError(ra_e.ValidationErrorException):
+    message = "Quota exceeded for resource '%(resource_name)s' in project %(project_id)s: %(current)s > %(limit)s"
+
+    def __init__(
+        self, resource_name: str, limit: int, current: int, project_id: sys_uuid.UUID
+    ):
+        super().__init__(
+            resource_name=resource_name,
+            limit=limit,
+            current=current,
+            project_id=project_id,
+        )
+        self.resource_name = resource_name
+        self.limit = limit
+        self.current = current
+        self.project_id = project_id
+
+
+# Built-in default quota limits keyed by resource_name (__tablename__).
+# These apply when no explicit QuotaLimit row exists for the project+resource.
+# A resource not listed here or in the DB has no quota limit (unlimited).
+DEFAULT_QUOTA_LIMIT = 1000
+DEFAULT_QUOTA_LIMITS: tp.Dict[str, int] = {
+    "net_lb": DEFAULT_QUOTA_LIMIT,
+    "compute_sets": DEFAULT_QUOTA_LIMIT,
+    "nodes": DEFAULT_QUOTA_LIMIT,
+    "secret_passwords": DEFAULT_QUOTA_LIMIT,
+    "secret_certificates": DEFAULT_QUOTA_LIMIT,
+    "secret_rsa_keys": DEFAULT_QUOTA_LIMIT,
+    "secret_ssh_keys": DEFAULT_QUOTA_LIMIT,
+}
+
+
+class QuotaModelMixin:
+    def _quota_limit(self, session) -> tp.Optional[int]:
+        limit = QuotaLimit.objects.get_one_or_none(
+            session=session,
+            filters={
+                "project_id": dm_filters.EQ(self.project_id),
+                "resource_name": dm_filters.EQ(self.__tablename__),
+            },
+        )
+        if limit is not None:
+            return limit.limit
+        return DEFAULT_QUOTA_LIMITS.get(self.__tablename__)
+
+    @staticmethod
+    def _quota_lock_key(project_id: sys_uuid.UUID, resource_name: str) -> int:
+        """Derive a positive bigint lock key for pg_advisory_xact_lock.
+
+        The key is a deterministic 63-bit hash of ``{project_id}:{resource_name}``,
+        so concurrent inserts for the same project+resource contend on the same
+        lock, while different project+resource pairs do not block each other.
+
+        Examples
+        --------
+        >>> import uuid
+        >>> pid = uuid.UUID("00000000-0000-0000-0000-000000000000")
+        >>> _quota_lock_key(pid, "net_lb")
+        8255477843906886675
+        >>> _quota_lock_key(pid, "compute_sets")
+        2792567584429316197
+        >>> # Different project → different key
+        >>> pid2 = uuid.UUID("12345678-1234-5678-1234-567812345678")
+        >>> _quota_lock_key(pid2, "net_lb")
+        1636646552482098955
+        """
+        key = f"{project_id}:{resource_name}"
+        digest = hashlib.sha256(key.encode()).digest()[:8]
+        return (
+            int.from_bytes(digest, byteorder="big", signed=False) & 0x7FFFFFFFFFFFFFFF
+        )
+
+    def _quota_acquire_lock(self, session: tp.Any) -> None:
+        lock_key = self._quota_lock_key(self.project_id, self.__tablename__)
+        session.execute("SELECT pg_advisory_xact_lock(%s)", [lock_key])
+
+    def _quota_reserve(self, session) -> None:
+        # Check limit, count existing reservations, create reservation.
+        # If limit exceeded, raises QuotaExceededError.
+        limit = self._quota_limit(session)
+        if limit is None:
+            return
+
+        self._quota_acquire_lock(session)
+
+        current = QuotaReservation.objects.count(
+            session=session,
+            filters={
+                "project_id": dm_filters.EQ(self.project_id),
+                "resource_name": dm_filters.EQ(self.__tablename__),
+            },
+        )
+
+        if current + 1 > limit:
+            raise QuotaExceededError(
+                resource_name=self.__tablename__,
+                limit=limit,
+                current=current + 1,
+                project_id=self.project_id,
+            )
+
+        reservation = QuotaReservation(
+            uuid=self.uuid,
+            project_id=self.project_id,
+            resource_name=self.__tablename__,
+        )
+        try:
+            reservation.insert(session=session)
+        except rs_e.ConflictRecords:
+            LOG.debug(
+                "Quota reservation already exists for %s",
+                self.uuid,
+            )
+
+    def _quota_release_all(self, session) -> None:
+        # Delete all reservations for this entity.
+        for reservation in QuotaReservation.objects.get_all(
+            filters={
+                "uuid": dm_filters.EQ(self.uuid),
+                "project_id": dm_filters.EQ(self.project_id),
+                "resource_name": dm_filters.EQ(self.__tablename__),
+            },
+            session=session,
+        ):
+            reservation.delete(session=session)
+
+    def insert(self, session=None):
+        # Reserve quota slot, then insert entity.
+        # The reservation is in the same transaction as the entity, so if the
+        # entity insert fails the entire transaction is rolled back and the
+        # reservation is cleaned up automatically — no explicit release needed.
+        if session is None:
+            engine = QuotaReservation._get_engine()
+            with engine.session_manager() as s:
+                self._quota_reserve(s)
+                super().insert(session=s)
+        else:
+            self._quota_reserve(session)
+            super().insert(session=session)
+
+    def delete(self, session=None):
+        # Release reservation, then delete entity.
+        if session is None:
+            engine = QuotaReservation._get_engine()
+            with engine.session_manager() as s:
+                self._quota_release_all(s)
+                super().delete(session=s)
+        else:
+            self._quota_release_all(session)
+            super().delete(session=session)
+
+
+def reconcile_quota_reservations(session):
+    cursor = session.execute(
+        "SELECT DISTINCT resource_name FROM quotas_reservations",
+        None,
+    )
+    resource_names = [row["resource_name"] for row in cursor]
+    for rname in resource_names:
+        session.execute(
+            "DELETE FROM quotas_reservations r "
+            "WHERE r.resource_name = %s "
+            "AND NOT EXISTS (SELECT 1 FROM " + rname + " e WHERE e.uuid = r.uuid)",
+            [rname],
+        )
+
+
+class QuotaReservation(
+    models.ModelWithUUID,
+    models.ModelWithTimestamp,
+    models.ModelWithProject,
+    orm.SQLStorableMixin,
+):
+    __tablename__ = "quotas_reservations"
+
+    resource_name = properties.property(
+        types.String(max_length=255),
+        required=True,
+        read_only=True,
+    )
+
+    @classmethod
+    def get_project_quota_summary(
+        cls,
+        session: tp.Optional[tp.Any] = None,
+        project_id: sys_uuid.UUID | str | None = None,
+    ) -> tp.List[tp.Dict[str, tp.Any]]:
+        """Get aggregated quota reservation information.
+
+        Returns a list of dictionaries with keys:
+          - project_id: UUID of the project
+          - resource_name: name of the resource
+          - reserved_count: number of reservations for this project+resource
+
+        If project_id is provided, filters to that specific project.
+        Otherwise, returns summaries for all projects.
+
+        If session is not provided, creates a temporary session using the default engine.
+        """
+
+        def req():
+            if project_id is not None:
+                return session.execute(
+                    "SELECT project_id, resource_name, COUNT(*) as reserved_count "
+                    "FROM quotas_reservations "
+                    "WHERE project_id = %s "
+                    "GROUP BY project_id, resource_name",
+                    [project_id],
+                )
+            else:
+                return session.execute(
+                    "SELECT project_id, resource_name, COUNT(*) as reserved_count "
+                    "FROM quotas_reservations "
+                    "GROUP BY project_id, resource_name",
+                    None,
+                )
+
+        if session is None:
+            engine = cls._get_engine()
+            with engine.session_manager() as session:
+                cursor = req()
+        else:
+            cursor = req()
+
+        result = [
+            {
+                "project_id": str(row["project_id"]),
+                "resource_name": row["resource_name"],
+                "reserved_count": row["reserved_count"],
+            }
+            for row in cursor
+        ]
+
+        return result
+
+
+class QuotaLimit(
+    models.ModelWithUUID,
+    models.ModelWithTimestamp,
+    models.ModelWithProject,
+    orm.SQLStorableMixin,
+):
+    __tablename__ = "quota_limits"
+
+    resource_name = properties.property(
+        types.String(max_length=255),
+        required=True,
+    )
+    limit = properties.property(
+        types.Integer(min_value=0),
+        required=True,
+    )