Add rate_limit option to Remote

https://pulp.plan.io/issues/7965
closes #7965

CHANGES/plugin_api/7965.feature

@@ -0,0 +1 @@
+Add ``rate_limit`` option to ``Remote``

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),
+        ),
+    ]

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.

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",
         )
 
 

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):

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)

requirements.txt

@@ -1,3 +1,4 @@
+asyncio-throttle~=1.0
 aiohttp~=3.7.2
 aiodns~=2.0.0
 aiofiles==0.6.0