openmedialibrary/oml/bandwidth.py

from time import time

import settings
import state
from websocket import trigger_event

import logging
logger = logging.getLogger(__name__)

class Bandwidth(object):
    _last = {}

    def __init__(self):
        self._upload = Bucket('uploadRate')
        self._download = Bucket('downloadRate')
        self.update()

    def update(self):
        bandwidth = {
            'up': self._upload.rate,
            'down': self._download.rate
        }
        if bandwidth != self._last:
            trigger_event('bandwidth', bandwidth)
            self._last = bandwidth
        state.main.call_later(1, self.update)

    def download(self, amount):
        return self._download.consume(amount)

    def upload(self, amount):
        return self._upload.consume(amount)


class Bucket(object):
    _rate = 0
    fill_rate = None

    def __init__(self, pref):
        self._pref = pref
        self.update()
        self.timestamp = time()
        self.rate_timestamp = time()

    def update(self):
        rate = settings.preferences.get(self._pref, '')
        if rate and rate.isdigit():
            rate = max(float(rate) * 1024, 0)
            if rate != self.fill_rate:
                self.capacity = max(rate, 2*16*1024)
                self._tokens = rate
                self.fill_rate = rate
        else:
            self.fill_rate = None

    def consume(self, tokens):
        self.update()
        if self.fill_rate is None:
            self._rate += tokens
            return True
        if tokens <= self.tokens:
            self._tokens -= tokens
            self._rate += tokens
        else:
            return False
        return True

    def get_tokens(self):
        now = time()
        if self._tokens < self.capacity:
            delta = self.fill_rate * (now - self.timestamp)
            self._tokens = min(self.capacity, self._tokens + delta)
        self.timestamp = now
        return self._tokens
    tokens = property(get_tokens)

    def get_rate(self):
        now = time()
        rate = self._rate * (now - self.rate_timestamp)
        self._rate = 0
        self.rate_timestamp = now
        return rate
    rate = property(get_rate)