SPACEDF-22039: [SpaceDF] - [BE] Refill free-plan device slot from deactivated devices on delete

apps/device/services/device_subscription_service.py

@@ -0,0 +1,99 @@
+from common.celery.task_senders import send_subscription_task
+from django.db import transaction
+
+from apps.device.models import Device
+
+
+def reconcile_device_limit(max_devices, deactivated_at=None):
+    rows = list(
+        Device.objects.values_list("id", "is_deactivated").order_by(
+            "created_at",
+            "id",
+        )
+    )
+    active_limit = len(rows) if max_devices is None else min(len(rows), max_devices)
+    active_ids = [device_id for device_id, _ in rows[:active_limit]]
+    suspended_ids = [device_id for device_id, _ in rows[active_limit:]]
+
+    active_id_set = set(active_ids)
+    suspended_id_set = set(suspended_ids)
+    reactivated_ids = [
+        device_id
+        for device_id, is_deactivated in rows
+        if is_deactivated and device_id in active_id_set
+    ]
+    deactivated_ids = [
+        device_id
+        for device_id, is_deactivated in rows
+        if not is_deactivated and device_id in suspended_id_set
+    ]
+
+    reactivated_count = (
+        Device.objects.filter(id__in=reactivated_ids, is_deactivated=True).update(
+            is_deactivated=False,
+            deactivated_at=None,
+        )
+        if reactivated_ids
+        else 0
+    )
+    deactivated_count = (
+        Device.objects.filter(id__in=deactivated_ids, is_deactivated=False).update(
+            is_deactivated=True,
+            deactivated_at=deactivated_at,
+        )
+        if deactivated_ids
+        else 0
+    )
+
+    return {
+        "total": len(rows),
+        "active_limit": active_limit,
+        "active_ids": active_ids,
+        "suspended_ids": suspended_ids,
+        "reactivated_ids": reactivated_ids,
+        "deactivated_ids": deactivated_ids,
+        "reactivated_count": reactivated_count,
+        "deactivated_count": deactivated_count,
+    }
+
+
+def reactivate_next_suspended_device():
+    device = (
+        Device.objects.filter(is_deactivated=True).order_by("created_at", "id").first()
+    )
+    if device is None:
+        return None
+
+    device.is_deactivated = False
+    device.deactivated_at = None
+    device.save(update_fields=["is_deactivated", "deactivated_at", "updated_at"])
+    return device.id
+
+
+def fill_active_slot_after_delete(fill_active_slot, org_slug):
+    if not fill_active_slot:
+        return
+
+    reactivated_id = reactivate_next_suspended_device()
+    if reactivated_id:
+        transaction.on_commit(
+            lambda: send_telemetry_subscription_task(
+                "upgrade",
+                [reactivated_id],
+                org_slug,
+            )
+        )
+
+
+def send_telemetry_subscription_task(lifecycle, device_ids, org_slug):
+    if not device_ids:
+        return
+    send_subscription_task(
+        service="telemetry",
+        lifecycle=lifecycle,
+        task_name=f"telemetry_{lifecycle}",
+        message={
+            "org_slug": org_slug,
+            "device_ids": [str(device_id) for device_id in device_ids],
+        },
+    )

apps/device/tasks.py

@@ -1,12 +1,14 @@
 import logging
 
 from common.apps.billing.constants import FeatureCode
-from common.celery.task_senders import send_subscription_task
 from common.celery.tasks import PermanentTaskError, task, tenant_shared_task
-from django.utils.dateparse import parse_datetime
 from django_tenants.utils import schema_context
 
 from apps.device.models import Device
+from apps.device.services.device_subscription_service import (
+    reconcile_device_limit,
+    send_telemetry_subscription_task,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -39,47 +41,13 @@ def device_downgrade_task(**kwargs):
     downgraded_at = kwargs.get("downgraded_at")
 
     with schema_context(org_slug):
-        devices = Device.objects.filter(is_deactivated=False).order_by("created_at")
-
-        excess_ids = list(devices.values_list("id", flat=True)[max_devices:])
-        count = (
-            Device.objects.filter(id__in=excess_ids, is_deactivated=False).update(
-                is_deactivated=True, deactivated_at=downgraded_at
-            )
-            if excess_ids
-            else 0
-        )
-        if not excess_ids and downgraded_at:
-            retry_timestamp = parse_datetime(downgraded_at)
-            if retry_timestamp is not None:
-                excess_ids = list(
-                    Device.objects.filter(
-                        is_deactivated=True,
-                        deactivated_at=retry_timestamp,
-                    ).values_list("id", flat=True)
-                )
-        if count:
-            logger.info(
-                "Downgrade: deactivated %s excess devices for org %s "
-                "(kept %s active out of %s total).",
-                count,
-                org_slug,
-                min(devices.count(), max_devices),
-                devices.count(),
-            )
+        result = reconcile_device_limit(max_devices, deactivated_at=downgraded_at)
+        count = result["deactivated_count"]
+        excess_ids = result["suspended_ids"]
 
     # Keep outside the schema_context - send_task is broker-only,
     # No DB access needed.
-    if excess_ids:
-        send_subscription_task(
-            service="telemetry",
-            lifecycle="downgrade",
-            task_name="telemetry_downgrade",
-            message={
-                "org_slug": org_slug,
-                "device_ids": [str(device_id) for device_id in excess_ids],
-            },
-        )
+    send_telemetry_subscription_task("downgrade", excess_ids, org_slug)
     return count
 
 
@@ -106,41 +74,12 @@ def device_upgrade_task(**kwargs):
         )
 
     with schema_context(org_slug):
-        if max_devices is None:
-            reactivated_ids = list(
-                Device.objects.order_by("created_at").values_list("id", flat=True)
-            )
-        else:
-            reactivated_ids = list(
-                Device.objects.order_by("created_at").values_list("id", flat=True)[
-                    :max_devices
-                ]
-            )
-        count = (
-            Device.objects.filter(id__in=reactivated_ids, is_deactivated=True).update(
-                is_deactivated=False, deactivated_at=None
-            )
-            if reactivated_ids
-            else 0
-        )
-        if count:
-            logger.info(
-                "Renewal: reactivated %s devices for org %s.",
-                count,
-                org_slug,
-            )
+        result = reconcile_device_limit(max_devices)
+        count = result["reactivated_count"]
+        reactivated_ids = result["active_ids"]
 
     # Cascade reactivation to telemetry entities.
-    if reactivated_ids:
-        send_subscription_task(
-            service="telemetry",
-            lifecycle="upgrade",
-            task_name="telemetry_upgrade",
-            message={
-                "org_slug": org_slug,
-                "device_ids": [str(device_id) for device_id in reactivated_ids],
-            },
-        )
+    send_telemetry_subscription_task("upgrade", reactivated_ids, org_slug)
     return count
 
 

apps/device/views.py

@@ -37,6 +37,9 @@
     UpdateSpaceDeviceSerializer,
 )
 from apps.device.services.device_profile_resolver import get_device_profile_context
+from apps.device.services.device_subscription_service import (
+    fill_active_slot_after_delete,
+)
 from apps.device.services.entity_properties_context import (
     _entity_properties_context,
     _organization_slug,
@@ -83,6 +86,13 @@ def get_serializer(self, *args, **kwargs):
             )
         return super().get_serializer(*args, **kwargs)
 
+    def perform_destroy(self, instance):
+        fill_active_slot = not instance.is_deactivated
+        org_slug = _organization_slug(self.request)
+        with transaction.atomic():
+            super().perform_destroy(instance)
+            fill_active_slot_after_delete(fill_active_slot, org_slug)
+
     @swagger_auto_schema(
         method="post",
         request_body=DeviceSerializer(many=True),