manyi/library/wechatpy/session/memcachedstorage.py

# -*- coding: utf-8 -*-
# !/usr/bin/env python

from __future__ import absolute_import, unicode_literals

import logging
import threading
import time

from library.wechatpy.session import SessionStorage
from library import to_text
from library.wechatpy.utils import json

logger = logging.getLogger(__name__)


class MemcachedStorage(SessionStorage):
    LOCAL_BUFFER = {}
    LOCAL_BUFFER_LOCK = threading.Lock()

    def __init__(self, mc, prefix = 'wechatpy'):
        for method_name in ('get', 'set', 'delete'):
            assert hasattr(mc, method_name)
        self.mc = mc
        self.prefix = prefix

    def key_name(self, key):
        return '{0}:{1}'.format(self.prefix, key)

    def get(self, key, default = None):
        key = self.key_name(key)

        mem_value = self.LOCAL_BUFFER.get(key, {})

        if mem_value and mem_value['expire'] > int(time.time()):
            logger.debug('=== WechatToken === hit memory. key = {}; expireAt = {}; value = {}'.format(
                    key, mem_value['expire'], mem_value['value']))
            return mem_value['value']
        else:
            with self.LOCAL_BUFFER_LOCK:
                mem_value = self.LOCAL_BUFFER.get(key, {})

                if mem_value and mem_value['expire'] > int(time.time()):
                    logger.debug(
                        '=== WechatToken === hit memory(double check). key = {}; expireAt = {}; value = {}'.format(
                            key, mem_value['expire'], mem_value['value']))
                    return mem_value['value']
                else:
                    logger.debug('=== WechatToken === miss memory. key = {}'.format(key))

                    value = self.mc.get(key, ignore_exc = False)

                    logger.debug("=== WechatToken === memcache: key = {}; result = {}".format(key, str(value)))

                    if value is None:
                        self.LOCAL_BUFFER.pop(key, None)
                        return default
                    else:
                        value = json.loads(to_text(value))
                        self.LOCAL_BUFFER[key] = value

                    logger.debug('=== WechatToken === set memory. key = {}; value = {};'.format(key, value))

                    return value['value']

    def set(self, key, value, ttl = 0):
        if value is None:
            return

        key = self.key_name(key)

        value = {
            'expire': int(time.time()) + ttl,
            'value': value
        }

        logger.debug("=== WechatToken === set memory: key = {}; value = {}; ttl = {}".format(key, value, ttl))

        with self.LOCAL_BUFFER_LOCK:
            self.LOCAL_BUFFER[key] = value
            logger.debug("=== WechatToken === set memcache: key = {}; value = {}; ttl = {}".format(key, value, ttl))
            self.mc.set(key, json.dumps(value), ttl)

    def delete(self, key):
        logger.debug("=== WechatToken === delete memcache: key = {};".format(key))

        key = self.key_name(key)
        self.mc.delete(key)