redis_lock.py

# coding=utf-8

import time
import socket
import os
from tornado import gen


# class _ReleasingContextManager(object):
#     def __init__(self, obj, future):
#         self._obj = obj
#         self._future = future
#
#     @gen.coroutine
#     def __enter__(self):
#         print '__enter__', self._future
#         yield self._future()
#
#     def __exit__(self, exc_type, exc_val, exc_tb):
#         print 'exit'
#         self._obj.release(force=True)


class AcquireLockTimeout(Exception):
    '''
    获取锁超时
    '''
    pass


class Lock(object):
    '''
    由于原有的系统采用阻塞版的redis, 实现锁的方式为 TimeoutTimer 去做阻塞操作
    '''
    LOCK_PREFIX = '_miaolegemi_lock_'
    LOCK_OWNER_PREFIX = '_miaolegemi_lock_owner_'

    def __init__(self, redis, key, timeout=None, ioloop=None):
        """
        :param redis: redis链接对象
        :type redis: cache.Cache
        :param key: 用于锁的key
        :type key: str
        :param timeout: 超时时间，时间为s
        :param timeout: float
        """
        self.lock_id = '%s_%d_%d' % (socket.gethostname(), os.getpid(), int(time.time() * 10000))
        self.redis = redis
        self.key = key
        self.timeout = timeout
        self._acquired = False
        assert isinstance(key, str)

    # @gen.coroutine
    # def acquire(self):
    #
    #     future = self._do_acquire
    #
    #     # if self.timeout:
    #     #     future = gen.with_timeout(timeout=self.timeout, future=future)
    #
    #     print 'will return'
    #     return _ReleasingContextManager(self, future)

    @gen.coroutine
    def acquire(self):
        print 'do_acquire'
        start_time = time.time()  # precision: ms

        while 1:
            val = self.redis.getset(self.LOCK_PREFIX + self.key, '1')

            if val and str(val) != '0':
                # this key is locked, we must wait
                # time_interval = 20ms
                if self.timeout and (time.time() - start_time) - self.timeout >= 0.001:  # tolerance 1ms
                    raise AcquireLockTimeout

                # now we wait
                yield gen.sleep(0.02)
            else:
                # we acquire the lock
                self.redis.set(self.LOCK_OWNER_PREFIX + self.key, self.lock_id)
                self._acquired = True
                print 'Get the lock'
                raise gen.Return(True)

    def release(self, force=False):
        '''
        :param force: 是否强制释放掉锁，并无视owner
        '''

        # if this lock never get lock
        if not self._acquired:
            return False

        # check if the key is locked
        is_lock, owner = self.redis.mget(self.LOCK_PREFIX + self.key,
                                         self.LOCK_OWNER_PREFIX + self.key)

        # return True means release successfully when
        # 1. key is lock and the lock instance has the lock
        # 2. key in not lock

        if is_lock and str(is_lock) != '0':
            # key is locked
            if owner != self.lock_id and not force:
                return False

            # the key is locked and we have the lock
            # release it
            print('delete')
            self.redis.delete(self.LOCK_PREFIX + self.key,
                              self.LOCK_OWNER_PREFIX + self.key)

        return True


if __name__ == '__main__':
    # test
    import redis
    import tornado.ioloop


    r = redis.StrictRedis()

    @gen.coroutine
    def test():
        print 'test'
        lock = Lock(r, 'test', timeout=0.5)

        print 'will get lock'
        try:
            yield lock.acquire()
        except gen.TimeoutError:
            print 'Timeout'
            return
        print '[***] critical area'
        # sleep
        yield gen.sleep(5)
        print 'awake!'
        lock.release()

        print 'done'



    ioloop = tornado.ioloop.IOLoop.current()


    ioloop.add_future(test(), lambda x: None)
    ioloop.add_future(test(), lambda x: None)
    ioloop.add_future(test(), lambda x: None)

    ioloop.start()