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

Author: boddumanohar

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

@@ -26,6 +26,7 @@
 from .....deployment import (
     AUTOSCALER_PVC_SUFFIX,
     STORAGE_PVC_SUFFIX,
+    WAL_IOPS_FRACTION,
     DeploymentParameters,
     ResizeParameters,
     branch_api_domain,
@@ -869,7 +870,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,
@@ -942,7 +944,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,
@@ -1006,6 +1009,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)
@@ -1022,7 +1026,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,
@@ -1711,6 +1716,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)
+
+    # Only split IOPS between data and WAL when a WAL volume exists (PITR enabled).
+    # Otherwise the full IOPS budget goes to the data volume.
+    # Catch only VelaKubernetesError (PVC/PV not found); other failures (e.g. missing
+    # CSI attributes) indicate corruption and must propagate.
+    try:
+        wal_volume, _ = await resolve_autoscaler_wal_volume_identifiers(namespace)
+    except VelaKubernetesError:
+        wal_volume = None
+
+    if wal_volume is not None:
+        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 wal_volume is not None:
+        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 volume not found for branch %s; skipping WAL IOPS update", branch_id)
 
 
 async def ensure_branch_storage_class(branch_id: Identifier, *, iops: int) -> str:
@@ -399,6 +426,7 @@ def _configure_vela_values(
     database_admin_password: str,
     pgbouncer_admin_password: str,
     storage_class_name: str,
+    wal_iops: int,
     use_existing_db_pvc: bool,
     pgbouncer_config: Mapping[str, int] | None,
     enable_file_storage: bool,
@@ -454,6 +482,7 @@ def _configure_vela_values(
     wal_persistence["create"] = not use_existing_db_pvc
     wal_persistence["size"] = PITR_WAL_PVC_SIZE
     wal_persistence["storageClassName"] = storage_class_name
+    wal_persistence.setdefault("annotations", {})["simplybk/qos-rw-iops"] = str(wal_iops)
     wal_persistence["claimName"] = wal_persistence.get("claimName") or (
         f"{_autoscaler_vm_name()}{AUTOSCALER_WAL_PVC_SUFFIX}"
     )
@@ -521,14 +550,21 @@ 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)
     values_content = _configure_vela_values(
         values_content,
         parameters=parameters,
         jwt_secret=jwt_secret,
         database_admin_password=database_admin_password,
         pgbouncer_admin_password=pgbouncer_admin_password,
         storage_class_name=storage_class_name,
+        wal_iops=wal_iops,
         use_existing_db_pvc=use_existing_db_pvc,
         pgbouncer_config=pgbouncer_config,
         enable_file_storage=parameters.enable_file_storage,

src/deployment/charts/vela/templates/autoscaler/wal-pvc.yaml

@@ -22,6 +22,12 @@ metadata:
   namespace: {{ .Release.Namespace }}
   labels:
     {{- include "vela.labels" . | nindent 4 }}
+  {{- if $pgWalPersistence.annotations }}
+  annotations:
+    {{- range $k, $v := $pgWalPersistence.annotations }}
+    {{ $k }}: {{ $v | quote }}
+    {{- end }}
+  {{- end }}
 spec:
   {{- if $pgWalPersistence.storageClassName }}
   storageClassName: {{ $pgWalPersistence.storageClassName }}