rlock.py

# -*- coding: utf-8 -*-
from __future__ import absolute_import

import logging

logger = logging.getLogger(__name__)

import time
import pytz
from datetime import datetime, timedelta

from django.utils.timezone import get_current_timezone, now
from django.conf import settings
import redis

r_conn = redis.Redis(host=getattr(settings, 'REDIS_HOST', 'localhost'),
                     port=getattr(settings, 'REDIS_PORT', 6379),
                     db=getattr(settings, 'REDIS_DB', 0),
                     password=getattr(settings, 'REDIS_PASSWORD', None)
)

def tz_aware(dt, timezone):
    """
    Makes a naive datetime.datetime in a given time zone aware.
    """
    if hasattr(timezone, 'localize'):
        # available for pytz time zones
        try:
            return timezone.localize(dt, is_dst=None)
        except (pytz.AmbiguousTimeError, pytz.NonExistentTimeError):
            return min(timezone.localize(dt, is_dst=True),
                       timezone.localize(dt, is_dst=False))
    else:
        # may be wrong around DST changes
        return dt.replace(tzinfo=timezone)

def dt_to_ts(datetime_obj):
    """Converts a datetime to epoch seconds"""
    epoch = tz_aware(datetime(1970, 1, 1),  get_current_timezone())
    return (datetime_obj - epoch).total_seconds()


def ts_to_dt(timestamp):
    return tz_aware(datetime.utcfromtimestamp(timestamp),
                    get_current_timezone())


class RLock(object):
    def __init__(self, r_key,
                 timeout=getattr(settings, 'REDIS_LOCK_TIMEOUT', 30)):
        self.timeout = timeout
        self.key = 'lock::%s' % r_key
        self.lock_time = None

    def acquire(self):
        while not self.lock_time:
            ts = time.time()
            if r_conn.setnx(self.key, ts + self.timeout):  # 1 if successful
                logger.debug('No previous lock. Acquired at %s', ts)
                self.lock_time = ts
            else:
                current_lock_ts = r_conn.get(self.key)
                if not current_lock_ts:
                    logger.debug('Lock expired before checking current ts. '
                                 'Re-attempting.')
                    # lock released between checks
                    continue
                else:
                    current_lock_ts = float(current_lock_ts)
                    if ts > current_lock_ts:
                        # stale lock - let's take it
                        logger.debug('Acquiring stale lock - expired at %s',
                                     current_lock_ts)
                        new_lock_ts = r_conn.getset(self.key, ts + self.timeout)
                        if not new_lock_ts:
                            # lock released between checks - we have the lock
                            logger.debug('Lock expired before confirming '
                                         'acquisition - acquired at %s', ts)
                            self.lock_time = ts
                        else:
                            new_lock_ts = float(new_lock_ts)
                            if new_lock_ts > ts:
                                logger.debug('Lock acquired by another thread. '
                                             'Re-attempting.')
                                # somebody else got the lock
                                continue
                            else:
                                # we got the lock
                                logger.debug('Lock acquired successfully at %s',
                                             ts)
                                self.lock_time = ts
                    else:
                        # lock is held by somebody else
                        logger.debug('Active lock held by another thread. '
                                     'Sleeping.')
                        time.sleep(getattr(settings,
                                           'REDIS_LOCK_WAIT_DELAY',
                                           0.1))
                        continue

    def release(self, *args):
        ts = time.time()
        if ts - self.timeout > self.lock_time:
            # lock expired - do nothing and let other clients handle it
            pass
        else:
            r_conn.delete(self.key)
        self.lock_time = None

    __enter__ = acquire
    __exit__ = release