diff --git a/CHANGES/plugin_api/7965.feature b/CHANGES/plugin_api/7965.feature new file mode 100644 index 0000000000..fa8055503d --- /dev/null +++ b/CHANGES/plugin_api/7965.feature @@ -0,0 +1 @@ +Add ``rate_limit`` option to ``Remote`` diff --git a/pulpcore/app/migrations/0056_remote_rate_limit.py b/pulpcore/app/migrations/0056_remote_rate_limit.py new file mode 100644 index 0000000000..ccc03e9bbd --- /dev/null +++ b/pulpcore/app/migrations/0056_remote_rate_limit.py @@ -0,0 +1,18 @@ +# Generated by Django 2.2.17 on 2021-01-27 18:49 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('core', '0055_label'), + ] + + operations = [ + migrations.AddField( + model_name='remote', + name='rate_limit', + field=models.IntegerField(null=True), + ), + ] diff --git a/pulpcore/app/models/repository.py b/pulpcore/app/models/repository.py index f1990a2d5c..a42c3c1365 100644 --- a/pulpcore/app/models/repository.py +++ b/pulpcore/app/models/repository.py @@ -7,6 +7,7 @@ import logging import django +from asyncio_throttle import Throttler from django.core.validators import MinValueValidator from django.db import models, transaction from django.urls import reverse @@ -215,6 +216,7 @@ class Remote(MasterModel): connect_timeout (models.FloatField): Value for aiohttp.ClientTimeout.connect sock_connect_timeout (models.FloatField): Value for aiohttp.ClientTimeout.sock_connect sock_read_timeout (models.FloatField): Value for aiohttp.ClientTimeout.sock_read + rate_limit (models.IntegerField): Limits total download rate in requests per second. """ TYPE = "remote" @@ -269,6 +271,7 @@ class Remote(MasterModel): null=True, validators=[MinValueValidator(0.0, "Timeout must be >= 0")] ) headers = JSONField(blank=True, null=True) + rate_limit = models.IntegerField(null=True) @property def download_factory(self): @@ -290,6 +293,26 @@ def download_factory(self): self._download_factory = DownloaderFactory(self) return self._download_factory + @property + def download_throttler(self): + """ + Return the Throttler which can be used to rate limit downloaders. + + Upon first access, the Throttler is instantiated and saved internally. + Plugin writers are expected to override when additional configuration of the + DownloaderFactory is needed. + + Returns: + Throttler: The instantiated Throttler to be used by get_downloader() + + """ + try: + return self._download_throttler + except AttributeError: + if self.rate_limit: + self._download_throttler = Throttler(rate_limit=self.rate_limit) + return self._download_throttler + def get_downloader(self, remote_artifact=None, url=None, **kwargs): """ Get a downloader from either a RemoteArtifact or URL that is configured with this Remote. diff --git a/pulpcore/app/serializers/repository.py b/pulpcore/app/serializers/repository.py index ad269bee6f..5fab7fb2f4 100644 --- a/pulpcore/app/serializers/repository.py +++ b/pulpcore/app/serializers/repository.py @@ -145,6 +145,11 @@ class RemoteSerializer(ModelSerializer): child=serializers.DictField(), help_text=_("Headers for aiohttp.Clientsession"), ) + rate_limit = serializers.IntegerField( + help_text=_("Limits total download rate in requests per second"), + allow_null=True, + required=False, + ) def validate_url(self, value): """ @@ -194,6 +199,7 @@ class Meta: "connect_timeout", "sock_connect_timeout", "sock_read_timeout", + "rate_limit", ) diff --git a/pulpcore/download/factory.py b/pulpcore/download/factory.py index da880b434c..cb3bd21eb4 100644 --- a/pulpcore/download/factory.py +++ b/pulpcore/download/factory.py @@ -188,6 +188,8 @@ class to be instantiated. login=self._remote.username, password=self._remote.password ) + kwargs["throttler"] = self._remote.download_throttler if self._remote.rate_limit else None + return download_class(url, **options, **kwargs) def _generic(self, download_class, url, **kwargs): diff --git a/pulpcore/download/http.py b/pulpcore/download/http.py index 7c3c45630f..ae4ea14855 100644 --- a/pulpcore/download/http.py +++ b/pulpcore/download/http.py @@ -120,6 +120,7 @@ def __init__( proxy_auth=None, headers_ready_callback=None, headers=None, + throttler=None, **kwargs, ): """ @@ -136,6 +137,7 @@ def __init__( available. The dictionary passed has the header names as the keys and header values as its values. e.g. `{'Transfer-Encoding': 'chunked'}` headers (dict): Headers to be submitted with the request. + throttler (asyncio_throttle.Throttler): Throttler for asyncio. kwargs (dict): This accepts the parameters of :class:`~pulpcore.plugin.download.BaseDownloader`. """ @@ -151,6 +153,7 @@ def __init__( self.proxy = proxy self.proxy_auth = proxy_auth self.headers_ready_callback = headers_ready_callback + self.download_throttler = throttler super().__init__(url, **kwargs) def raise_for_status(self, response): @@ -208,6 +211,8 @@ async def _run(self, extra_data=None): Args: extra_data (dict): Extra data passed by the downloader. """ + if self.download_throttler: + await self.download_throttler.acquire() async with self.session.get(self.url, proxy=self.proxy, auth=self.auth) as response: self.raise_for_status(response) to_return = await self._handle_response(response) diff --git a/requirements.txt b/requirements.txt index e42b41fe0c..ce50c2d6ec 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +asyncio-throttle~=1.0 aiohttp~=3.7.2 aiodns~=2.0.0 aiofiles==0.6.0