root/branches/crazycache/dejavu/storage/memcache.py

"""Memcached client for use with storememcached.MemcachedStorageManager."""

try:
    import cPickle as pickle
except ImportError:
    import pickle
import socket
import threading
import time


class Flagset(object):
    """A set of coercion rules for memcache data encoding."""
    
    def push(self, value):
        """Return (flags, coerced value) for output."""
        if isinstance(value, basestring):
            return 0, value
        elif isinstance(value, int):
            return 2, str(value)
        elif isinstance(value, long):
            return 4, str(value)
        return 1, pickle.dumps(value, 2)
    
    def pull(self, value, flags):
        """Return value, coerced for input according to the given flags."""
        if flags == 0:
            return value
        elif flags & 2:
            return int(value)
        elif flags & 4:
            return long(value)
        elif flags & 1:
            try:
                return pickle.loads(value)
            except:
                return None
        raise ValueError("unknown flags %x for value %r" % (flags, value))


class Client(object):
    """Memcache client interface.
    
    servers: a sequence of "host[:port]" strings.
    """
    
    flagset = Flagset()
    
    def __init__(self, servers):
        self.servers = [Server(s) for s in servers]
    
    def _key_to_server(self, key):
        """returns a server without locking it"""
        return self.servers[hash(key) % len(self.servers)]
    
    def _get_server(self, key):
        """returns a server and locks it"""
        server = self._key_to_server(key)
        if server.acquire():
            return server
        return None
    
    def delete(self, key, time=0):
        """Delete the cached value for the given key."""
        if time:
            cmd = "delete %s %d" % (key, time)
        else:
            cmd = "delete %s" % key
        server = self._get_server(key)
        if server:
            try:
                r = server.store(cmd)
                if r not in ('DELETED', 'NOT_FOUND'):
                    raise IOError(r)
            finally:
                server.release()
    
    def _set(self, cmd, key, val, time):
        server = self._get_server(key)
        if server:
            try:
                flags, val = self.flagset.push(val)
                msg = ("%s %s %d %d %d\r\n%s" %
                       (cmd, key, flags, time, len(val), val))
                r = server.store(msg)
                if r != "STORED":
                    raise IOError(r)
            finally:
                server.release()
    
    def add(self, key, val, time=0):
        """Add a new key/value pair to the cache (unless it exists)."""
        self._set("add", key, val, time)
    
    def replace(self, key, val, time=0):
        """Replace an existing key/value pair, only if it exists."""
        self._set("replace", key, val, time)
    
    def set(self, key, val, time=0):
        """Set the given key/value pair."""
        self._set("set", key, val, time)
    
    def get(self, key):
        """Retrieve the cached value for the given key."""
        server = self._get_server(key)
        if server:
            try:
                server.retrieve("get %s" % key)
                line = server.readline()
                if not line.startswith('VALUE'):
                    return None
                
                _, _, flags, length = line.split()
                buf = server.recv(int(length) + 2)[:-2]
                
                if server.readline() != "END":
                    return None
                return self.flagset.pull(buf, int(flags))
            finally:
                server.release()
    
    def get_multi(self, keys):
        """Return a dict of cached values for the given keys.
        
        The returned dict will contain only those entries which
        were found in the cache.
        """
        
        # the list of memcached servers we need
        servers = {}
        for key in keys:
            server = self._key_to_server(key)
            if server not in servers:
                servers[server] = [str(key)]
            else:
                servers[server].append(str(key))
        
        results = {}
        try:
            # Connect to all servers
            for server in servers:
                if not server.acquire():
                    # server couldn't be acquired (no socket), remove it
                    del servers[server]
                    continue
            
            # request the keys in every server
            for server, lkeys in servers.items():
                try:
                    server.retrieve("get %s" % " ".join(lkeys))
                except socket.error:
                    # remove the dead server
                    del servers[server]
                    server.release()
            
            # Read the results from each server
            for server in servers:
                try:
                    line = server.readline()
                    while line and line.strip() != 'END':
                        _, key, flags, length = line.split()
                        buf = server.recv(int(length) + 2)[:-2]
                        results[key] = self.flagset.pull(buf, int(flags))
                        line = server.readline()
                except socket.error:
                    pass
        finally:
            for server in servers:
                server.release()
        
        return results
    
    def flush_all(self):
        """Expire all data currently in the memcache servers."""
        for s in self.servers:
            if s.acquire():
                try:
                    s.flush_all()
                finally:
                    s.release()
    
    def disconnect_all(self):
        """Disconnect from all servers."""
        for s in self.servers:
            s.close()


class Server(object):
    
    deaduntil = 0
    socket = None
    retry_delay = 5
    
    def __init__(self, host):
        if ":" in host:
            self.ip, port = host.split(":")
            self.port = int(port)
        else:
            self.ip, self.port = host, 11211
        self._socket_lock = threading.Lock()
    
    def mark_dead(self):
        self.deaduntil = time.time() + self.retry_delay
        self.close()
    
    def acquire(self):
        if self.deaduntil and self.deaduntil > time.time():
            return None
        else:
            self.deaduntil = 0
        
        if not self.socket:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(1.0)
            try:
                s.connect((self.ip, self.port))
            except (socket.error, socket.timeout), msg:
                # For #120
                try:
                    detail = msg[1]
                except IndexError:
                    detail = repr(msg)
                self.mark_dead()
                return None
            self.socket = s
        
        self._socket_lock.acquire()
        return self.socket
    
    def release(self):
        """Release a connection."""
        self._socket_lock.release()
    
    def close(self):
        if self.socket:
            self.socket.close()
            self.socket = None
    
    def store(self, command):
        """Send the given command string and return its single-line response."""
        try:
            self.socket.sendall(command + '\r\n')
        except socket.error, exc:
            self.mark_dead()
            raise
        return self.readline()
    
    def retrieve(self, command):
        """Send the given command string and let the caller recv the response."""
        try:
            self.socket.sendall(command + '\r\n')
        except socket.error, exc:
            self.mark_dead()
            raise
    
    def readline(self):
        try:
            buffers = ''
            recv = self.socket.recv
            while 1:
                data = recv(1)
                if not data:
                    self.mark_dead()
                    break
                if data == '\n' and buffers and buffers[-1] == '\r':
                    return(buffers[:-1])
                buffers = buffers + data
            return buffers
        except socket.error, exc:
            self.mark_dead()
            raise
    
    def recv(self, size):
        try:
            buf = ''
            buflen = 0
            recv = self.socket.recv
            while True:
                buf += recv(size - buflen)
                buflen = len(buf)
                if buflen >= size:
                    break
            if buflen != size:
                raise IOError("%d bytes received (expected %d)." %
                              (buflen, size), buf)
            if not buf.endswith('\r\n'):
                raise IOError("Response did not end with CRLF.", buf)
            
            return buf
        except socket.error, exc:
            self.mark_dead()
            raise
    
    def flush_all(self):
        """Expire all data currently in the server."""
        r = self.store('flush_all')
        if r != "OK":
            raise IOError(r)
    
    def __repr__(self):
        d = ''
        if self.deaduntil:
            d = " (dead until %d)" % self.deaduntil
        return "<memcached at %s port: %d%s>" % (self.ip, self.port, d)