Реализация кеша с ограничением по числу элементов на Python — два решения: простое и посложнее

Формулировка задачи

Предположим, что у нас есть необходимость иметь некий сервис, который бы отдавал нам по запросу какую-либо информацию, и отдавал как можно быстрее. Что для этого делает любой нормальный человек? Налаживает кэширование наиболее часто запрашиваемых данных. При этом, если хоть немного задуматься о перспективе, то размеры кэша необходимо ограничивать.
Для простоты реализации в случае Питона сделаем ограничение по числу элементов в кэше (здесь предполагается, что данные более-менее однородны по размеру, а также учитывается специфика, что определить объём памяти, реально занимаемый каким-либо Питоновским объектом — весьма нетривиальная задача, кому интересно, пусть пожалует сюда), а для того, чтобы кэш содержал как можно более часто используемую информацию — построим его по принципу least recently used, т.е. чем более давно запрашивали кусочек информации, тем больше у него шансов «вылететь» из кэша.

О двух решениях (попроще и посложнее) я и расскажу в этой статье.

Философское отступление

Я обратил внимание, что зачастую написанный на Питоне код не блещет оптимизациями по потреблению памяти или по скорости (здесь можно вскользь заметить, что это не только вина программистов-пользователей языка, хорошего инструментария просто нет, или по крайней мере я о таком не знаю (да, я в курсе про cProfile, но он помогает далеко не всегда, например, в нём нет attach-detach; искать же утечки памяти — вообще занятие не для слабонервных, вот если pympler допилят...), если кто подскажет — буду благодарен). В основном это даже правильно, так как обычно Питон применяют в случае, когда время исполнения программы или потребление памяти не критично, зато критично время написания кода и простота его дальнейшей поддержки, поэтому очевидное, пусть и менее экономичное решение будет удобнее.
Хотя вышесказанное совсем не означает, что любой код на Питоне нужно писать «абы как». Кроме того, иногда производительность и потребление памяти могут стать критичными, даже если код крутится на серверной машине с хорошим процессором и тоннами памяти.
Как раз такой случай (нехватка CPU time) и подвиг меня на исследование этого вопроса и замену одного кэширующего класса другим.

Простое решение

Часто говорят — «простое решение — самое верное», хотя в случае с программированием это далеко не всегда так, не правда ли? Однако, если есть задача обеспечения кэша, но нет особых требований к скорости (например, то, что использует этот кэш, само по себе очень медленное, и нет смысла вкладываться в сложную реализацию), то можно обойтись и тривиальными решениями.

Реализация

Итак, относительно простое решение:

import weakref
class SimpleCache:
    def __init__(self, cacheSize):
        self.cache = weakref.WeakValueDictionary()
        self.cache_LRU = []
        self.cacheSize = cacheSize

    def __getitem__(self, key):
        # no exception catching - let calling code handle them if any
        value = self.cache[key]
        return self.touchCache(key, value)

    def __setitem__(self, key, value):
        self.touchCache(key, value)

    def touchCache(self, key, value):
        self.cache[key] = value
        if value in self.cache_LRU:
            self.cache_LRU.remove(value)
        self.cache_LRU = self.cache_LRU[-(self.cacheSize-1):] + [ value ]
        return value

Использование

Работать с таким кэшем после создания можно, как с обычным dict-ом, но при чтении/записи элемента он будет помещён в конец очереди, а за счёт комбинации WeakValueDictionary() и списка, который хранит не более cacheSize элементов, мы получаем ограничение по количеству сохранённых данных.
Таким образом, после исполнения куска кода

cache = SimpleCache(2)
cache[1] = 'a'
cache[2] = 'b'
cache[3] = 'c'

в кэше останутся только записи 'b' и 'c' ('a', как самая старая, будет вытеснена).

Преимущества и недостатки

Преимущество — относительная простота (около 20 строк кода, после прочтения документации по WeakValueDictionary принцип работы становится понятен).
Недостаток — скорость работы, т.к. при каждом «касании» кэша приходится пересоздавать список целиком, удаляя из произвольного его места элемент (на самом деле там происходит аж целая куча длинных операций по работе со списком, так, «if value in self.cache_LRU» — линейный поиск по списку, потом .remove() — ещё раз линейный поиск, потом ещё берётся срез — т.е. создаётся почти полная копия списка). Словом, относительно долго.

Сложное решение

Теперь вдумаемся — а как можно ускорить такой класс? Очевидно, основные проблемы у нас именно с поддержанием cache_LRU в актуальном состоянии — как я уже говорил, поиск по нему, последующее удаление элемента и пересоздание — дорогие операции. Тут нам на помощь придёт такая вещь, как двунаправленный связный список — он обеспечит нам поддержку «последний использованный — в конце», ну а WeakValueDictionary() поможет с быстрым поиском по этому списку (поиск по словарю идёт куда быстрее линейного перебора, поскольку внутри Питоновские словари реализуют что-то вроде бинарных деревьев по хешам ключей (тут Остапа понесло — могу честно сказать, что в исходники не смотрел, поэтому только предполагаю, как именно устроен поиск по словарю)

Реализация

Для начала создаём класс — элемент списка, в который будем оборачивать хранимые данные:

    class Element(object):
        __slots__ = ['prev', 'next', 'value', '__init__', '__weakref__']
        def __init__(self, value):
            self.prev, self.next, self.value = None, None, value

Тут надо заметить использование такой штуки, как __slots__, позволяющей существенно сэкономить память и немного — производительность по сравнению с такой же реализацией класса, но без этого атрибута (что это такое и с чем его едят — в официальной документации).
Теперь создаём класс кэша (класс элемента засунем внутрь «во избежание»):

import weakref
class FastCache:
    class Element(object):
        __slots__ = ['prev', 'next', 'value', '__init__', '__weakref__']
        def __init__(self, value):
            self.prev, self.next, self.value = None, None, value

    def __init__(self, maxCount):
        self.dict = weakref.WeakValueDictionary()
        self.head = None
        self.tail = None
        self.count = 0
        self.maxCount = maxCount

    def _removeElement(self, element):
        prev, next = element.prev, element.next
        if prev:
            assert prev.next == element
            prev.next = next
        elif self.head == element:
            self.head = next

        if next:
            assert next.prev == element
            next.prev = prev
        elif self.tail == element:
            self.tail = prev
        element.prev, element.next = None, None
        assert self.count >= 1
        self.count -= 1

    def _appendElement(self, element):
        if element is None:
            return
        element.prev, element.next = self.tail, None
        if self.head is None:
            self.head = element
        if self.tail is not None:
            self.tail.next = element
        self.tail = element
        self.count += 1

    def get(self, key, *arg):
        element = self.dict.get(key, None)
        if element:
            self._removeElement(element)
            self._appendElement(element)
            return element.value
        elif len(*arg):
            return arg[0]
        else:
            raise KeyError("'%s' is not found in the dictionary", str(key))

    def __len__(self):
        return len(self.dict)

    def __getitem__(self, key):
        element = self.dict[key]
        # At this point the element is not None
        self._removeElement(element)
        self._appendElement(element)
        return element.value

    def __setitem__(self, key, value):
        try:
            element = self.dict[key]
            self._removeElement(element)
        except KeyError:
            if self.count == self.maxCount:
                self._removeElement(self.head)
        element = FastCache.Element(value)
        self._appendElement(element)
        self.dict[key] = element
        
    def __del__(self):
        while self.head:
            self._removeElement(self.head)

Тут можно обратить внимание на следующие существенные/интересные моменты:

• реализация метода get() — слегка отличается от стандартной для словарей, т.к. если не задано значение по умолчанию, а ключ отсутствует, то выкидывает исключение (работает так же, как cache[key]), а если есть значение по умолчанию, то возвращает его
• наличие метода __del__ обязательно, в противном случае получаем (или можем получить) утечку всего кэша — ведь все элементы списка друг на друга ссылаются, значит, их счётчики ссылок никогда не обнулятся без посторонней помощи; кстати, как выяснилось, встроенный (по крайней мере в 2.6) garbage collector хоть и вроде как собирает простейшие циклы ссылок, но с этим списком не справляется
• при желании можно слегка поднять производительность, выкинув assert'ы

Использование

Полностью аналогично предыдущей реализации (да и логика внутри похожа, только использует не простой, встроенный в Питон тип list, а реализует свой двунаправленный список с шахматами и поэтессами).

Сравнение

Голословные утверждения — это одно, а беспристрастные цифры — совсем другое. Поэтому я не призываю верить мне на слово, наоборот — измерим эти классы!
Тут нам на помощь приходит встроенный в Питон модуль timeit:

class StubClass:
    def __init__(self, something):
        self.something = something

def testCache(cacheClass, cacheSize, repeat):
    d = cacheClass(cacheSize)
    for i in xrange(repeat * cacheSize):
        d[i] = StubClass(i)

def minMaxAvg(lst):
    return min(lst), max(lst), 1.0 * sum(lst) / len(lst)
    
if __name__ == '__main__':
    import gc, sys, pdb
    gc.set_debug(gc.DEBUG_LEAK)
    BIG_NUM, REPEAT1, REPEAT2, REPEAT3 = 100, 10, 10, 10
    
    import timeit
    T1 = timeit.Timer(lambda:testCache(SimpleCache, BIG_NUM, REPEAT1))
    T2 = timeit.Timer(lambda:testCache(FastCache, BIG_NUM, REPEAT1))
    
    tt1, tt2 = [x.repeat(REPEAT2, REPEAT3) for x in (T1, T2)]
    print 'Simple: min %.5f, max %.5f, avg %.5f' % minMaxAvg(tt1)
    print 'Fast  : min %.5f, max %.5f, avg %.5f' % minMaxAvg(tt2)

Результаты запуска на моей машине (Intel Core i5 2540M 2.6GHz, ActivePython 2.7.2 x64-bit):

Simple: min 2.23016, max 2.38483, avg 2.31895
Fast  : min 0.09513, max 0.10594, avg 0.09916

Разница довольно ощутима — порядка 20 раз!

По потреблению памяти — запускаем предыдущий скрипт с другой секцией main и смотрим потребление памяти интерпретатором Питона через диспетчер задач:

if __name__ == '__main__':
    import time
    print 'measure me - no cache'
    try:
        while True:
            time.sleep(10)
    except:
        # let user interrupt this with Ctrl-C
        pass
    testCache(FastCache, 1000, 1)
    print 'measure me - cached'
    while True:
        time.sleep(10)
    exit()

Результаты измерений — в таблице:

Как видим — сложное решение не только заметно быстрее (в ~20 раз) добавляет элементы, но и потребляет немного меньше памяти, чем простое.

В той конкретной задаче, которую я решал, создавая такой кэш, его замена позволила сократить время отдачи запроса с 90 секунд до 70 примерно (более плотная профилировка и переписывание почти всей логики генерации ответа потом помогла довести время ответа до 30 секунд, но это уже совсем другая история).

Автор: JustAMan