webapp/codereview/memcache.py

# Copyright 2008 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility support for dealing with memcache."""

from google.appengine.api import memcache
from google.appengine.runtime import DeadlineExceededError

import cStringIO
import pickle
import zlib


class Key(object):
  def __init__(self, key, timeout=None, compress=False):
    if compress:
      key += '(z)'
    self._key = key
    self._timeout = timeout
    self._compress = compress

  def get(self, compute=None):
    try:
      r = memcache.get(self._key)
      if r and self._compress:
        r = pickle.load(cStringIO.StringIO(zlib.decompress(r)))
    except DeadlineExceededError:
      r = None

    if r is not None:
      return r

    if compute is not None:
      r = compute()
      if r is not None:
        self.set(r)
    return r

  def clear(self):
    try:
      memcache.delete(self._key)
    except DeadlineExceededError:
      pass

  def set(self, value):
    try:
      if self._compress:
        buf = cStringIO.StringIO()
        pickle.dump(value, buf, -1)
        value = zlib.compress(buf.getvalue())

      if self._timeout is None:
        memcache.set(self._key, value)
      else:
        memcache.set(self._key, value, self._timeout)
    except DeadlineExceededError:
      pass


class CachedDict(object):
  """A cache of zero or more memcache keys.  The dictionary of
     acquired keys is also cached locally in this process, but
     can be forcibly cleared by clear_local.
  """
  def __init__(self,
               prefix,
               timeout,
               compute_one = None,
               compute_multi = None):
    self._prefix = prefix
    self._timeout = timeout
    self._cache = {}
    self._prefetch = set()

    if compute_multi:
      self._compute = compute_multi
    elif compute_one:
      self._compute = lambda x: map(compute_one, x)
    else:
      raise ValueError, 'compute_one or compute_multi required'

  def prefetch(self, keys):
    for item_key in keys:
      if item_key not in self._cache:
        self._prefetch.add(item_key)

  def get_multi(self, keys):
    result = {}

    not_local = []
    for item_key in keys:
      try:
        result[item_key] = self._cache[item_key]
      except KeyError:
        not_local.append(item_key)

    if not_local:
      to_get = []
      to_get.extend(not_local)
      to_get.extend(self._prefetch)
      self._prefetch = set()

      try:
        cached = memcache.get_multi(to_get, self._prefix)
      except DeadlineExceededError:
        cached = {}

      to_compute = []
      for item_key in to_get:
        try:
          r = cached[item_key]
        except KeyError:
          to_compute.append(item_key)
          continue
        self._cache[item_key] = r
        result[item_key] = r

      if to_compute:
        to_cache = {}
        for item_key, r in zip(to_compute, self._compute(to_compute)):
          if r is not None:
            to_cache[item_key] = r
          self._cache[item_key] = r
          result[item_key] = r

        if to_cache:
          try:
            memcache.set_multi(to_cache, self._timeout, self._prefix)
          except DeadlineExceededError:
            pass
    return result

  def get(self, key):
    return get_multi([key])[key]

  def clear_local(self):
    self._cache = {}
    self._prefetch = set()

  def clear(self):
    try:
      memcache.delete_multi(self._prefix)
    except DeadlineExceededError:
      pass
    self.clear_local()