WIP: rate limitsrate-limits

4 files changed, 227 insertions, 6 deletions

diff --git a/pandora/client.py b/pandora/client.py
index a98056c..ff27678 100644
--- a/pandora/client.py
+++ b/pandora/client.py
@@ -14,6 +14,7 @@ For simplicity use a client builder from pandora.clientbuilder to create an
 instance of a client.
 """
 from . import errors
+from .ratelimit import WarningTokenBucket
 
 
 class BaseAPIClient:
@@ -36,7 +37,27 @@ class BaseAPIClient:
         partner_password,
         device,
         default_audio_quality=MED_AUDIO_QUALITY,
+        rate_limiter=WarningTokenBucket,
     ):
+        """Initialize an API Client
+
+        transport
+            instance of a Pandora transport
+        partner_user
+            partner username
+        partner_password
+            partner password
+        device
+            device type identifier
+        default_audio_quality
+            audio quality level, one of the *_AUDIO_QUALITY constants in the
+            BaseAPIClient class
+        rate_limiter
+            class (not instance) implementing the pandora.ratelimit.TokenBucket
+            interface. Used by various client components to handle rate
+            limits with the Pandora API. The default rate limited warns when
+            the rate limit has been exceeded but does not enforce the limit.
+        """
         self.transport = transport
         self.partner_user = partner_user
         self.partner_password = partner_password
@@ -45,6 +66,17 @@ class BaseAPIClient:
         self.username = None
         self.password = None
 
+        # Global rate limiter for all methods, allows a call rate of 2 calls
+        # per second. This limit is based on nothing but seems sane to prevent
+        # runaway code from hitting Pandora too hard.
+        self._api_limiter = rate_limiter(120, 1, 1)
+
+        # Rate limiter for the get_playlist API which has a much lower
+        # server-side limit than other APIs. This was determined
+        # experimentally. This is applied before and in addition to the global
+        # API rate limit.
+        self._playlist_limiter = rate_limiter(5, 1, 1)
+
     def _partner_login(self):
         partner = self.transport(
             "auth.partnerLogin",
@@ -98,6 +130,8 @@ class BaseAPIClient:
                 return []
 
     def __call__(self, method, **kwargs):
+        self._api_limiter.consume(1)
+
         try:
             return self.transport(method, **kwargs)
         except errors.InvalidAuthToken:
@@ -125,6 +159,8 @@ class APIClient(BaseAPIClient):
     def get_playlist(self, station_token, additional_urls=None):
         from .models.playlist import Playlist
 
+        self._playlist_limiter.consume(1)
+
         if additional_urls is None:
             additional_urls = []
 
diff --git a/pandora/clientbuilder.py b/pandora/clientbuilder.py
index 2b11ffb..1eb3811 100644
--- a/pandora/clientbuilder.py
+++ b/pandora/clientbuilder.py
@@ -8,6 +8,7 @@ import os.path
 
 from configparser import ConfigParser
 from .client import APIClient
+from .ratelimit import WarningTokenBucket
 from .transport import Encryptor, APITransport, DEFAULT_API_HOST
 
 
@@ -94,8 +95,9 @@ class APIClientBuilder:
 
     DEFAULT_CLIENT_CLASS = APIClient
 
-    def __init__(self, client_class=None):
+    def __init__(self, client_class=None, rate_limiter=WarningTokenBucket):
         self.client_class = client_class or self.DEFAULT_CLIENT_CLASS
+        self.rate_limiter = rate_limiter
 
     def build_from_settings_dict(self, settings):
         enc = Encryptor(settings["DECRYPTION_KEY"], settings["ENCRYPTION_KEY"])
@@ -116,6 +118,7 @@ class APIClientBuilder:
             settings["PARTNER_PASSWORD"],
             settings["DEVICE"],
             quality,
+            self.rate_limiter,
         )
 
 
diff --git a/pandora/ratelimit.py b/pandora/ratelimit.py
new file mode 100644
index 0000000..427ede9
--- /dev/null
+++ b/pandora/ratelimit.py
@@ -0,0 +1,164 @@
+"""
+Pandora Rate Limiter
+"""
+import time
+import warnings
+
+from .errors import PandoraException
+
+
+class RateLimitExceeded(PandoraException):
+    """Exception thrown when rate limit is exceeded
+    """
+
+    code = 0
+    message = "Rate Limit Exceeded"
+
+
+class TokenBucketCallbacks(object):
+    """Interface for TokenBucket Callbacks
+
+    These methods get called by the token bucket during certain parts of the
+    bucket lifecycle.
+    """
+
+    def near_depletion(self, tokens_left):
+        """Bucket near depletion callback
+
+        This callback is called when the token bucket is nearly depleted, as
+        defined by the depletion_level member of the TokenBucket class. This
+        method provides a hook for clients to warn when the token bucket is
+        nearly consumed. The return value of this method is ignored.
+        """
+        return
+
+    def depleted(self, tokens_wanted, tokens_left):
+        """Bucket depletion callback
+
+        This callback is called when the token bucket has been depleted.
+        Returns a boolean indicating if the bucket should call its default
+        depletion handler. If this method returns False the token bucket will
+        not return any tokens but will also not call the default token
+        depletion handler.
+
+        This method is useful if clients want to customize depletion behavior.
+        """
+        return True
+
+    def consumed(self, tokens_wanted, tokens_left):
+        """Bucket consumption callback
+
+        This callback is called when a client successfully consumes tokens from
+        a bucket. The return value of this method is ignored.
+        """
+        return
+
+
+class TokenBucket(object):
+    def __init__(self, capacity, refill_tokens, refill_rate, callbacks=None):
+        """Initialize a rate limiter
+
+        capacity
+            the number of tokens the bucket holds when completely full
+        refill_tokens
+            number of tokens to add to the bucket during each refill cycle
+        refill_rate
+            the number of seconds between token refills
+        """
+        self.capacity = capacity
+        self.refill_tokens = refill_tokens
+        self.refill_rate = refill_rate
+        self.callbacks = callbacks or TokenBucketCallbacks()
+
+        # Depletion level at which the near_depletion callback is called.
+        # Defaults to 20% of the bucket available. Not exposed in the
+        # initializer because this should generally be good enough.
+        self.depletion_level = self.capacity / 5
+
+        self._available_tokens = capacity
+        self._last_refill = time.time()
+
+    @classmethod
+    def creator(cls, callbacks=None):
+        """Returns a TokenBucket creator
+
+        This method is used when clients want to customize the callbacks but
+        defer class construction to the real consumer.
+        """
+
+        def constructor(capacity, refill_tokens, refill_rate):
+            return cls(capacity, refill_tokens, refill_rate, callbacks)
+
+        return constructor
+
+    def _replentish(self):
+        now = time.time()
+
+        refill_unit = round((now - self._last_refill) / self.refill_rate)
+        if refill_unit > 0:
+            self._last_refill = now
+            self._available_tokens = min(
+                self.capacity, refill_unit * self.refill_tokens
+            )
+
+    def _insufficient_tokens(self, tokens_wanted):
+        return
+
+    def _sufficient_tokens(self, tokens_wanted):
+        return
+
+    def consume(self, tokens=1):
+        """Consume a number of tokens from the bucket
+
+        May return a boolean indicating that sufficient tokens were available
+        for consumption. Implementations may have different behaviour when the
+        bucket is empty. This method may block or throw.
+        """
+        self._replentish()
+
+        if self._available_tokens >= tokens:
+            self._available_tokens -= tokens
+
+            if self._available_tokens <= self.depletion_level:
+                self.callbacks.near_depletion(self._available_tokens)
+
+            self.callbacks.consumed(tokens, self._available_tokens)
+            self._sufficient_tokens(tokens)
+
+            return True
+        else:
+            if self.callbacks.depleted(tokens, self._available_tokens):
+                self._insufficient_tokens(tokens)
+            return False
+
+
+class BlockingTokenBucket(TokenBucket):
+    """Token bucket that blocks on exhaustion
+    """
+
+    def _insufficient_tokens(self, tokens_wanted):
+        tokens_per_second = self.refill_rate / self.refill_tokens
+        excess_tokens = tokens_wanted - self._available_tokens
+        time.sleep(tokens_per_second * excess_tokens)
+
+
+class ThrowingTokenBucket(TokenBucket):
+    """Token bucket that throws on exhaustion
+    """
+
+    def _insufficient_tokens(self, tokens_wanted):
+        raise RateLimitExceeded("Unable to acquire enough tokens")
+
+
+class WarningTokenBucket(TokenBucket):
+    """Token bucket that warns on exhaustion
+
+    This token bucket doesn't enforce the rate limit and is designed to be a
+    softer policy for backwards compatability with code that was written before
+    rate limiting was added.
+    """
+
+    def _insufficient_tokens(self, tokens_wanted):
+        warnings.warn(
+            "Pandora API rate limit exceeded!", warnings.ResourceWarning
+        )
diff --git a/pydora/player.py b/pydora/player.py
index c4ce2e0..d2b1388 100644
--- a/pydora/player.py
+++ b/pydora/player.py
@@ -10,7 +10,9 @@ import os
 import sys
 import logging
 import argparse
-from pandora import clientbuilder
+from pandora.clientbuilder import PydoraConfigFileBuilder
+from pandora.clientbuilder import PianobarConfigFileBuilder
+from pandora.ratelimit import BlockingTokenBucket, TokenBucketCallbacks
 
 from .utils import Colors, Screen
 from .audio_backend import RemoteVLC
@@ -46,6 +48,18 @@ class PlayerCallbacks:
         pass
 
 
+class PlayerTokenBucketCallbacks(TokenBucketCallbacks):
+    def __init__(self, screen):
+        self.screen = screen
+
+    def near_depletion(self, tokens_left):
+        self.screen.print_error("Near Pandora API call rate limit. Slow down!")
+
+    def depleted(self, tokens_wanted, tokens_left):
+        self.screen.print_error("Pandora API call rate exceeded. Please wait!")
+        return True
+
+
 class PlayerApp:
 
     CMD_MAP = {
@@ -75,10 +89,10 @@ class PlayerApp:
             try:
                 host, port = vlc_net.split(":")
                 player = RemoteVLC(host, port, self, sys.stdin)
-                Screen.print_success("Using Remote VLC")
+                self.screen.print_success("Using Remote VLC")
                 return player
             except PlayerUnusable:
-                Screen.print_error("Unable to connect to vlc")
+                self.screen.print_error("Unable to connect to vlc")
                 raise
 
         try:
@@ -99,12 +113,16 @@ class PlayerApp:
         sys.exit(1)
 
     def get_client(self):
+        tb_maker = BlockingTokenBucket.creator(
+            PlayerTokenBucketCallbacks(self.screen)
+        )
+
         cfg_file = os.environ.get("PYDORA_CFG", "")
-        builder = clientbuilder.PydoraConfigFileBuilder(cfg_file)
+        builder = PydoraConfigFileBuilder(cfg_file, rate_limiter=tb_maker)
         if builder.file_exists:
             return builder.build()
 
-        builder = clientbuilder.PianobarConfigFileBuilder()
+        builder = PianobarConfigFileBuilder(rate_limiter=tb_maker)
         if builder.file_exists:
             return builder.build()