IFC-1530 Add weighted pool resource generic (#6511)

This change adds weight to resources of pools in order to have a more
precise control over the order of resources during the allocation
process.

Author: gmazoyer

backend/infrahub/core/constants/infrahubkind.py

@@ -70,3 +70,4 @@
 USERVALIDATOR = "CoreUserValidator"
 VALIDATOR = "CoreValidator"
 WEBHOOK = "CoreWebhook"
+WEIGHTED_POOL_RESOURCE = "CoreWeightedPoolResource"

backend/infrahub/core/node/resource_manager/ip_address_pool.py

@@ -81,11 +81,15 @@ async def get_resource(
         return node
 
     async def get_next(self, db: InfrahubDatabase, prefixlen: int | None = None) -> IPAddressType:
-        # Measure utilization of all prefixes identified as resources
         resources = await self.resources.get_peers(db=db)  # type: ignore[attr-defined]
         ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
 
-        for resource in resources.values():
+        try:
+            weighted_resources = sorted(resources.values(), key=lambda r: r.allocation_weight.value or 0, reverse=True)
+        except AttributeError:
+            weighted_resources = list(resources.values())
+
+        for resource in weighted_resources:
             ip_prefix = ipaddress.ip_network(resource.prefix.value)  # type: ignore[attr-defined]
             prefix_length = prefixlen or ip_prefix.prefixlen
 

backend/infrahub/core/node/resource_manager/ip_prefix_pool.py

@@ -88,11 +88,15 @@ async def get_resource(
         return node
 
     async def get_next(self, db: InfrahubDatabase, prefixlen: int) -> IPNetworkType:
-        # Measure utilization of all prefixes identified as resources
         resources = await self.resources.get_peers(db=db)  # type: ignore[attr-defined]
         ip_namespace = await self.ip_namespace.get_peer(db=db)  # type: ignore[attr-defined]
 
-        for resource in resources.values():
+        try:
+            weighted_resources = sorted(resources.values(), key=lambda r: r.allocation_weight.value or 0, reverse=True)
+        except AttributeError:
+            weighted_resources = list(resources.values())
+
+        for resource in weighted_resources:
             subnets = await get_subnets(
                 db=db,
                 ip_prefix=ipaddress.ip_network(resource.prefix.value),  # type: ignore[attr-defined]

backend/infrahub/core/protocols.py

@@ -209,6 +209,10 @@ class CoreWebhook(CoreNode):
     validate_certificates: BooleanOptional
 
 
+class CoreWeightedPoolResource(CoreNode):
+    allocation_weight: IntegerOptional
+
+
 class LineageOwner(CoreNode):
     pass
 

backend/infrahub/core/schema/definitions/core/__init__.py

@@ -56,7 +56,13 @@
     core_user_validator,
 )
 from .repository import core_generic_repository, core_read_only_repository, core_repository
-from .resource_pool import core_ip_address_pool, core_ip_prefix_pool, core_number_pool, core_resource_pool
+from .resource_pool import (
+    core_ip_address_pool,
+    core_ip_prefix_pool,
+    core_number_pool,
+    core_resource_pool,
+    core_weighted_pool_resource,
+)
 from .template import core_object_component_template, core_object_template
 from .transform import core_transform, core_transform_jinja2, core_transform_python
 from .webhook import core_custom_webhook, core_standard_webhook, core_webhook
@@ -81,6 +87,7 @@
         builtin_ip_prefix,
         builtin_ip_address,
         core_resource_pool,
+        core_weighted_pool_resource,
         core_generic_account,
         core_base_permission,
         core_credential,

backend/infrahub/core/schema/definitions/core/resource_pool.py

@@ -32,6 +32,17 @@
     ],
 )
 
+core_weighted_pool_resource = GenericSchema(
+    name="WeightedPoolResource",
+    namespace="Core",
+    label="Weighted Pool Resource",
+    description="Resource to be used in a pool, its weight is used to determine its priority on allocation.",
+    include_in_menu=False,
+    branch=BranchSupportType.AWARE,
+    generate_profile=False,
+    attributes=[Attr(name="allocation_weight", label="Weight", kind="Number", optional=True, order_weight=10000)],
+)
+
 core_ip_prefix_pool = NodeSchema(
     name="IPPrefixPool",
     namespace="Core",

backend/tests/unit/conftest.py

@@ -2505,7 +2505,7 @@ async def ipam_schema() -> SchemaRoot:
                 "order_by": ["prefix__value"],
                 "display_labels": ["prefix__value"],
                 "branch": BranchSupportType.AWARE.value,
-                "inherit_from": [InfrahubKind.IPPREFIX],
+                "inherit_from": [InfrahubKind.IPPREFIX, InfrahubKind.WEIGHTED_POOL_RESOURCE],
             },
             {
                 "name": "IPAddress",

backend/tests/unit/core/resource_manager/test_ipaddress_pool.py

@@ -42,6 +42,46 @@ async def test_get_next(
     assert next_prefix.id == next_prefix2.id
 
 
+async def test_get_next_weighted(
+    db: InfrahubDatabase,
+    default_branch: Branch,
+    default_ipnamespace: Node,
+    register_ipam_schema: SchemaBranch,
+    ip_dataset_prefix_v4,
+):
+    ns1 = ip_dataset_prefix_v4["ns1"]
+    net144 = ip_dataset_prefix_v4["net144"]
+    net145 = ip_dataset_prefix_v4["net145"]
+
+    net144.allocation_weight.value = 100
+    await net144.save(db=db)
+    net145.allocation_weight.value = 200
+    await net145.save(db=db)
+
+    adress_pool_schema = registry.schema.get_node_schema(name=InfrahubKind.IPADDRESSPOOL, branch=default_branch)
+
+    pool = await CoreIPAddressPool.init(schema=adress_pool_schema, db=db)
+    await pool.new(
+        db=db, name="pool1", resources=[net144, net145], ip_namespace=ns1, default_address_type="IpamIPAddress"
+    )
+    await pool.save(db=db)
+
+    assert pool
+
+    next_address = await pool.get_next(db=db, prefixlen=30)
+    assert str(next_address) == "10.10.3.2/30"
+
+    next_prefix = await pool.get_resource(
+        db=db, address_type="IpamIPAddress", identifier="item1", branch=default_branch
+    )
+    assert next_prefix
+
+    next_prefix2 = await pool.get_resource(
+        db=db, address_type="IpamIPAddress", identifier="item1", branch=default_branch
+    )
+    assert next_prefix.id == next_prefix2.id
+
+
 async def test_get_next_full(
     db: InfrahubDatabase,
     default_branch: Branch,

backend/tests/unit/core/resource_manager/test_prefix_pool.py

@@ -43,6 +43,44 @@ async def test_get_next(
     assert next_prefix.id == next_prefix2.id
 
 
+async def test_get_next_weighted(
+    db: InfrahubDatabase,
+    default_branch: Branch,
+    default_ipnamespace: Node,
+    register_ipam_schema: SchemaBranch,
+    ip_dataset_prefix_v4,
+):
+    ns1 = ip_dataset_prefix_v4["ns1"]
+    net140 = ip_dataset_prefix_v4["net140"]
+    net141 = ip_dataset_prefix_v4["net141"]
+
+    net140.allocation_weight.value = 100
+    await net140.save(db=db)
+    net141.allocation_weight.value = 200
+    await net141.save(db=db)
+
+    prefix_pool_schema = registry.schema.get_node_schema(name=InfrahubKind.IPPREFIXPOOL, branch=default_branch)
+
+    pool = await CoreIPPrefixPool.init(schema=prefix_pool_schema, db=db)
+    await pool.new(db=db, name="pool1", resources=[net140, net141], ip_namespace=ns1)
+    await pool.save(db=db)
+
+    assert pool
+
+    next_subnet = await pool.get_next(db=db, prefixlen=17)
+    assert str(next_subnet) == "10.11.0.0/17"
+
+    next_prefix = await pool.get_resource(
+        db=db, prefixlen=17, prefix_type="IpamIPPrefix", member_type="prefix", identifier="item1", branch=default_branch
+    )
+    assert next_prefix
+
+    next_prefix2 = await pool.get_resource(
+        db=db, prefixlen=17, prefix_type="IpamIPPrefix", member_type="prefix", identifier="item1", branch=default_branch
+    )
+    assert next_prefix.id == next_prefix2.id
+
+
 async def test_get_one(
     db: InfrahubDatabase,
     default_branch: Branch,

backend/tests/unit/message_bus/operations/requests/test_proposed_change.py

@@ -140,7 +140,7 @@ async def test_get_proposed_change_schema_integrity_constraints(
     )
     non_generate_profile_constraints = [c for c in constraints if c.constraint_name != "node.generate_profile.update"]
     # should be updated/removed when ConstraintValidatorDeterminer is updated (#2592)
-    assert len(constraints) == 184
+    assert len(constraints) == 185
     assert len(non_generate_profile_constraints) == 110
     dumped_constraints = [c.model_dump() for c in non_generate_profile_constraints]
     assert {