iops split for PITR branches: 25% to WAL and 75% to DATA

Author: boddumanohar

src/api/organization/project/branch/__init__.py

@@ -25,6 +25,7 @@
 from ....._util.crypto import encrypt_with_passphrase, generate_keys
 from .....database import AsyncSessionLocal, SessionDep
 from .....deployment import (
+    WAL_IOPS_FRACTION,
     DeploymentParameters,
     ResizeParameters,
     branch_api_domain,
@@ -754,7 +755,8 @@ async def _clone_branch_environment_task(
     storage_class_name: str | None = None
     if copy_data:
         try:
-            storage_class_name = await ensure_branch_storage_class(branch_id, iops=parameters.iops)
+            data_iops = max(1, round(parameters.iops * (1 - WAL_IOPS_FRACTION))) if pitr_enabled else parameters.iops
+            storage_class_name = await ensure_branch_storage_class(branch_id, iops=data_iops)
             await clone_branch_database_volume(
                 source_branch_id=source_branch_id,
                 target_branch_id=branch_id,
@@ -829,7 +831,8 @@ async def _restore_branch_environment_task(
     await _persist_branch_status(branch_id, BranchServiceStatus.CREATING)
     storage_class_name: str | None = None
     try:
-        storage_class_name = await ensure_branch_storage_class(branch_id, iops=parameters.iops)
+        data_iops = max(1, round(parameters.iops * (1 - WAL_IOPS_FRACTION))) if pitr_enabled else parameters.iops
+        storage_class_name = await ensure_branch_storage_class(branch_id, iops=data_iops)
         await restore_branch_database_volume_from_snapshot(
             source_branch_id=source_branch_id,
             target_branch_id=branch_id,
@@ -894,6 +897,7 @@ async def _restore_branch_environment_in_place_task(
     snapshot_namespace: str,
     snapshot_name: str,
     snapshot_content_name: str | None,
+    pitr_enabled: bool,
     recovery_target_time: datetime | None = None,
 ) -> None:
     await _persist_branch_status(branch_id, BranchServiceStatus.RESTARTING)
@@ -910,7 +914,8 @@ async def _restore_branch_environment_in_place_task(
             return
 
     try:
-        storage_class_name = await ensure_branch_storage_class(branch_id, iops=parameters.iops)
+        data_iops = max(1, round(parameters.iops * (1 - WAL_IOPS_FRACTION))) if pitr_enabled else parameters.iops
+        storage_class_name = await ensure_branch_storage_class(branch_id, iops=data_iops)
         await restore_branch_database_volume_from_snapshot(
             source_branch_id=source_branch_id,
             target_branch_id=branch_id,
@@ -1600,6 +1605,7 @@ async def restore(
             snapshot_namespace=snapshot_namespace,
             snapshot_name=snapshot_name,
             snapshot_content_name=snapshot_content_name,
+            pitr_enabled=branch.pitr_enabled,
             recovery_target_time=restore_target if isinstance(restore_target, datetime) else None,
         )
     )

src/deployment/__init__.py

@@ -79,6 +79,7 @@
 AUTOSCALER_PVC_SUFFIX = "-block-data"
 AUTOSCALER_WAL_PVC_SUFFIX = "-pg-wal"
 PITR_WAL_PVC_SIZE = "100Gi"
+WAL_IOPS_FRACTION = 0.25
 _LOAD_BALANCER_TIMEOUT_SECONDS = float(600)
 _LOAD_BALANCER_POLL_INTERVAL_SECONDS = float(2)
 _OVERLAY_IP_TIMEOUT_SECONDS = float(300)
@@ -356,13 +357,39 @@ async def update_branch_volume_iops(branch_id: Identifier, iops: int) -> None:
     namespace = deployment_namespace(branch_id)
 
     volume, _ = await resolve_autoscaler_volume_identifiers(namespace)
+
+    # Detect PITR by checking whether the WAL PVC exists.
+    wal_pvc_name = f"{_autoscaler_vm_name()}{AUTOSCALER_WAL_PVC_SUFFIX}"
+    try:
+        await kube_service.get_persistent_volume_claim(namespace, wal_pvc_name)
+        pitr_active = True
+    except VelaKubernetesError:
+        pitr_active = False
+
+    if pitr_active:
+        data_iops = max(1, round(iops * (1 - WAL_IOPS_FRACTION)))
+        wal_iops = max(1, round(iops * WAL_IOPS_FRACTION))
+    else:
+        data_iops = iops
+
     try:
         async with create_simplyblock_api() as sb_api:
-            await sb_api.update_volume(volume=volume, payload={"max_rw_iops": iops})
+            await sb_api.update_volume(volume=volume, payload={"max_rw_iops": data_iops})
     except VelaSimplyblockAPIError as exc:
         raise VelaDeploymentError("Failed to update volume") from exc
 
-    logger.info("Updated Simplyblock volume %s IOPS to %s", volume, iops)
+    logger.info("Updated Simplyblock data volume %s IOPS to %s", volume, data_iops)
+
+    if pitr_active:
+        wal_volume, _ = await resolve_autoscaler_wal_volume_identifiers(namespace)
+        try:
+            async with create_simplyblock_api() as sb_api:
+                await sb_api.update_volume(volume=wal_volume, payload={"max_rw_iops": wal_iops})
+        except VelaSimplyblockAPIError as exc:
+            raise VelaDeploymentError("Failed to update WAL volume") from exc
+        logger.info("Updated Simplyblock WAL volume %s IOPS to %s", wal_volume, wal_iops)
+    else:
+        logger.info("WAL PVC absent for branch %s; full IOPS budget applied to data volume", branch_id)
 
 
 async def ensure_branch_storage_class(branch_id: Identifier, *, iops: int) -> str:
@@ -397,15 +424,18 @@ async def _create_fresh_pvcs(
     parameters: DeploymentParameters,
     *,
     pitr_enabled: bool,
+    wal_iops: int = 0,
 ) -> None:
     """Create empty block PVCs for a brand-new deployment before helm install."""
     from .kubernetes.pvc import create_pvc
 
     access_modes = ["ReadWriteMany"]
 
-    def _pvc(name: str, size: str) -> kubernetes_client.V1PersistentVolumeClaim:
+    def _pvc(
+        name: str, size: str, annotations: dict[str, str] | None = None
+    ) -> kubernetes_client.V1PersistentVolumeClaim:
         return kubernetes_client.V1PersistentVolumeClaim(
-            metadata=kubernetes_client.V1ObjectMeta(name=name),
+            metadata=kubernetes_client.V1ObjectMeta(name=name, annotations=annotations),
             spec=kubernetes_client.V1PersistentVolumeClaimSpec(
                 access_modes=access_modes,
                 storage_class_name=storage_class_name,
@@ -419,7 +449,8 @@ def _pvc(name: str, size: str) -> kubernetes_client.V1PersistentVolumeClaim:
 
     if pitr_enabled:
         wal_pvc_name = f"{_autoscaler_vm_name()}{AUTOSCALER_WAL_PVC_SUFFIX}"
-        await create_pvc(namespace, _pvc(wal_pvc_name, PITR_WAL_PVC_SIZE))
+        wal_annotations = {"simplybk/qos-rw-iops": str(wal_iops)} if wal_iops > 0 else None
+        await create_pvc(namespace, _pvc(wal_pvc_name, PITR_WAL_PVC_SIZE, annotations=wal_annotations))
 
 
 def _configure_vela_values(
@@ -548,10 +579,18 @@ async def create_vela_config(
     postgresql_resource = resources.files(__package__).joinpath("postgresql.conf")
     values_content = _load_chart_values(chart)
 
-    storage_class_name = await ensure_branch_storage_class(branch_id, iops=parameters.iops)
+    if pitr_enabled:
+        data_iops = max(1, round(parameters.iops * (1 - WAL_IOPS_FRACTION)))
+        wal_iops = max(1, round(parameters.iops * WAL_IOPS_FRACTION))
+    else:
+        data_iops = parameters.iops
+        wal_iops = 0
+    storage_class_name = await ensure_branch_storage_class(branch_id, iops=data_iops)
 
     if not use_existing_db_pvc:
-        await _create_fresh_pvcs(namespace, storage_class_name, parameters, pitr_enabled=pitr_enabled)
+        await _create_fresh_pvcs(
+            namespace, storage_class_name, parameters, pitr_enabled=pitr_enabled, wal_iops=wal_iops
+        )
 
     values_content = _configure_vela_values(
         values_content,