-
Notifications
You must be signed in to change notification settings - Fork 30
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: make subscription license provisioning asynchronous and batched
ENT-8269
- Loading branch information
Showing
8 changed files
with
344 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -85,6 +85,9 @@ private.py | |
# Editor Temp Files | ||
*.swp | ||
|
||
# pyenv | ||
.python-version | ||
|
||
*.trace | ||
|
||
docs/_build/ | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
23 changes: 23 additions & 0 deletions
23
license_manager/apps/subscriptions/migrations/0064_subscriptionplan_desired_num_licenses.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,23 @@ | ||
# Generated by Django 4.2.9 on 2024-01-24 03:39 | ||
|
||
from django.db import migrations, models | ||
|
||
|
||
class Migration(migrations.Migration): | ||
|
||
dependencies = [ | ||
('subscriptions', '0063_transfer_all_licenses'), | ||
] | ||
|
||
operations = [ | ||
migrations.AddField( | ||
model_name='historicalsubscriptionplan', | ||
name='desired_num_licenses', | ||
field=models.PositiveIntegerField(blank=True, help_text='Total number of licenses that should exist for this SubscriptionPlan. The total license count (provisioned asynchronously) will reach the desired amount eventually. Empty (NULL) means no attempts will be made to asynchronously provision licenses.', null=True, verbose_name='Desired Number of Licenses'), | ||
), | ||
migrations.AddField( | ||
model_name='subscriptionplan', | ||
name='desired_num_licenses', | ||
field=models.PositiveIntegerField(blank=True, help_text='Total number of licenses that should exist for this SubscriptionPlan. The total license count (provisioned asynchronously) will reach the desired amount eventually. Empty (NULL) means no attempts will be made to asynchronously provision licenses.', null=True, verbose_name='Desired Number of Licenses'), | ||
), | ||
] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,133 @@ | ||
""" | ||
Celery tasks for the subscriptions app. | ||
""" | ||
import functools | ||
import logging | ||
from datetime import timedelta | ||
|
||
from celery import shared_task, states | ||
from celery_utils.logged_task import LoggedTask | ||
from django.db import IntegrityError | ||
from django.db.utils import OperationalError | ||
from django_celery_results.models import TaskResult | ||
|
||
from license_manager.apps.subscriptions.models import SubscriptionPlan | ||
from license_manager.apps.subscriptions.utils import ( | ||
batch_counts, | ||
localized_utcnow, | ||
) | ||
|
||
|
||
logger = logging.getLogger(__name__) | ||
|
||
ONE_HOUR = timedelta(hours=1) | ||
TASK_RETRY_SECONDS = 60 | ||
PROVISION_LICENSES_BATCH_SIZE = 300 | ||
|
||
|
||
class RequiredTaskUnreadyError(Exception): | ||
""" | ||
An exception representing a state where one type of task that is required | ||
to be complete before another task is run is not in a ready state. | ||
""" | ||
|
||
|
||
def unready_tasks(celery_task, time_delta): | ||
""" | ||
Returns any unready tasks with the name of the given celery task | ||
that were created within the given (now - time_delta, now) range. | ||
The unready celery states are | ||
{'RECEIVED', 'REJECTED', 'STARTED', 'PENDING', 'RETRY'}. | ||
https://docs.celeryproject.org/en/v5.0.5/reference/celery.states.html#unready-states | ||
Args: | ||
celery_task: A celery task definition or "type" (not an applied task "instance"), | ||
for example, ``update_catalog_metadata_task``. | ||
time_delta: A datetime.timedelta indicating how for back to look for unready tasks of this type. | ||
""" | ||
return TaskResult.objects.filter( | ||
task_name=celery_task.name, | ||
date_created__gte=localized_utcnow() - time_delta, | ||
status__in=states.UNREADY_STATES, | ||
) | ||
|
||
|
||
def already_running_semaphore(time_delta=None): | ||
""" | ||
Celery Task decorator that wraps a bound (bind=True) task. If another task with the same (name, args, kwargs) as | ||
the given task was executed in the time between `time_delta` and now, and the task still has not completed, defer | ||
running this task (by retrying until all other tasks are completed). | ||
`time_delta` defaults to one hour. | ||
Args: | ||
time_delta (datetime.timedelta): An optional timedelta that specifies how far back | ||
to look for the same task. | ||
""" | ||
def decorator(task): | ||
@functools.wraps(task) | ||
def wrapped_task(self, *args, **kwargs): | ||
delta = time_delta or ONE_HOUR | ||
if unready_tasks(self, delta).exists(): | ||
logger.info( | ||
f'Deferring task {self.name} with id {self.request.id} ' | ||
f'and args: {self.request.args}, kwargs: {self.request.kwargs}, ' | ||
'since another task run with the same arguments has not yet completed.' | ||
) | ||
raise self.retry(exc=RequiredTaskUnreadyError()) | ||
return task(self, *args, **kwargs) | ||
return wrapped_task | ||
return decorator | ||
|
||
|
||
class LoggedTaskWithRetry(LoggedTask): # pylint: disable=abstract-method | ||
""" | ||
Shared base task that allows tasks that raise some common exceptions to retry automatically. | ||
See https://docs.celeryproject.org/en/stable/userguide/tasks.html#automatic-retry-for-known-exceptions for | ||
more documentation. | ||
""" | ||
autoretry_for = ( | ||
IntegrityError, | ||
OperationalError, | ||
) | ||
retry_kwargs = {'max_retries': 5} | ||
# Use exponential backoff for retrying tasks | ||
retry_backoff = True | ||
# Add randomness to backoff delays to prevent all tasks in queue from executing simultaneously | ||
retry_jitter = True | ||
|
||
|
||
@shared_task(base=LoggedTaskWithRetry, bind=True, default_retry_delay=TASK_RETRY_SECONDS) | ||
@already_running_semaphore() | ||
def provision_licenses_task(self, subscription_plan_uuid=None): # pylint: disable=unused-argument | ||
""" | ||
For a given subscription plan, try to make its count of licenses match the number defined by the | ||
`desired_num_licenses` field of that subscription plan. Never decrease the count of licenses; if there are already | ||
more licenses than `desired_num_licenses`, do nothing. | ||
Args: | ||
subscription_plan_uuid (str): UUID of the SubscriptionPlan object to provision licenses for. | ||
""" | ||
subscription_plan = SubscriptionPlan.objects.get(uuid=subscription_plan_uuid) | ||
if not subscription_plan.desired_num_licenses: | ||
logger.info( | ||
f'Skipping task {self.name} with id {self.request.id} ' | ||
f'and args: {self.request.args}, kwargs: {self.request.kwargs}, ' | ||
f'because desired_num_licenses is not set on this subscription plan.' | ||
) | ||
return | ||
license_count_gap = subscription_plan.desired_num_licenses - subscription_plan.num_licenses | ||
if license_count_gap <= 0: | ||
logger.info( | ||
f'Skipping task {self.name} with id {self.request.id} ' | ||
f'and args: {self.request.args}, kwargs: {self.request.kwargs}, ' | ||
f'because the actual license count ({subscription_plan.num_licenses}) ' | ||
f'already meets or exceeds the desired license count ({subscription_plan.desired_num_licenses}).' | ||
) | ||
return | ||
|
||
# There's work to do, creating licenses! It should be safe to not re-check the license count between batches | ||
# because we lock this subscription plan anyway (via @already_running_semaphore decorator). | ||
for batch_count in batch_counts(license_count_gap, batch_size=PROVISION_LICENSES_BATCH_SIZE): | ||
subscription_plan.increase_num_licenses(batch_count) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,99 @@ | ||
""" | ||
Tests for subscriptions app celery tasks | ||
""" | ||
from datetime import datetime, timedelta | ||
from unittest import mock | ||
from uuid import uuid4 | ||
|
||
import ddt | ||
import freezegun | ||
import pytest | ||
from braze.exceptions import BrazeClientError | ||
from django.conf import settings | ||
from django.test import TestCase | ||
from django.test.utils import override_settings | ||
from freezegun import freeze_time | ||
from requests import models | ||
|
||
from license_manager.apps.subscriptions import tasks | ||
from license_manager.apps.subscriptions.models import License, SubscriptionPlan | ||
from license_manager.apps.subscriptions.tests.factories import ( | ||
LicenseFactory, | ||
SubscriptionPlanFactory, | ||
) | ||
|
||
|
||
# pylint: disable=unused-argument | ||
@ddt.ddt | ||
class ProvisionLicensesTaskTests(TestCase): | ||
""" | ||
Tests for provision_licenses_task. | ||
""" | ||
def setUp(self): | ||
super().setUp() | ||
self.subscription_plan = SubscriptionPlanFactory() | ||
|
||
# For all cases below, assume batch size of 5. | ||
@ddt.data( | ||
# Don't add licenses if none are desired. | ||
{ | ||
'num_initial_licenses': 0, | ||
'desired_num_licenses': 0, | ||
'expected_num_licenses': 0, | ||
}, | ||
# Create fewer licenses than one batch. | ||
{ | ||
'num_initial_licenses': 0, | ||
'desired_num_licenses': 1, | ||
'expected_num_licenses': 1, | ||
}, | ||
# Create licenses that span multiple batches. | ||
{ | ||
'num_initial_licenses': 0, | ||
'desired_num_licenses': 6, | ||
'expected_num_licenses': 6, | ||
}, | ||
# Don't add more licenses if the goal has already been reached. | ||
{ | ||
'num_initial_licenses': 10, | ||
'desired_num_licenses': 10, | ||
'expected_num_licenses': 10, | ||
}, | ||
# Create fewer licenses than one batch (starting with 10 initially). | ||
{ | ||
'num_initial_licenses': 10, | ||
'desired_num_licenses': 11, | ||
'expected_num_licenses': 11, | ||
}, | ||
# Create licenses that span multiple batches (starting with 10 initially). | ||
{ | ||
'num_initial_licenses': 10, | ||
'desired_num_licenses': 16, | ||
'expected_num_licenses': 16, | ||
}, | ||
# Don't remove licenses if the desired number of licenses is smaller than count of existing licenses. | ||
{ | ||
'num_initial_licenses': 20, | ||
'desired_num_licenses': 10, | ||
'expected_num_licenses': 20, | ||
}, | ||
# Don't add or remove licenses if the desired number of licenses is None. | ||
{ | ||
'num_initial_licenses': 20, | ||
'desired_num_licenses': None, | ||
'expected_num_licenses': 20, | ||
}, | ||
) | ||
@ddt.unpack | ||
@mock.patch('license_manager.apps.subscriptions.tasks.PROVISION_LICENSES_BATCH_SIZE', 5) | ||
def test_provision_licenses_task(self, num_initial_licenses, desired_num_licenses, expected_num_licenses): | ||
""" | ||
Test provision_licenses_task. | ||
""" | ||
self.subscription_plan.desired_num_licenses = desired_num_licenses | ||
self.subscription_plan.save() | ||
self.subscription_plan.increase_num_licenses(num_initial_licenses) | ||
|
||
tasks.provision_licenses_task(subscription_plan_uuid=self.subscription_plan.uuid) | ||
|
||
assert self.subscription_plan.num_licenses == expected_num_licenses |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters