class TestRedlock(unittest.TestCase): def setUp(self): self.redlock = Redlock([{"host": "localhost"}]) def test_lock(self): lock = self.redlock.lock("pants", 100) self.assertEqual(lock.resource, "pants") self.redlock.unlock(lock) lock = self.redlock.lock("pants", 10) self.redlock.unlock(lock) def test_blocked(self): lock = self.redlock.lock("pants", 1000) bad = self.redlock.lock("pants", 10) self.assertFalse(bad) self.redlock.unlock(lock) def test_bad_connection_info(self): with self.assertRaises(Warning): Redlock([{"cat": "hog"}]) def test_py3_compatible_encoding(self): lock = self.redlock.lock("pants", 1000) key = self.redlock.servers[0].get("pants") self.assertEquals(lock.key, key)
class TestRedlock(APITestCase): def setUp(self): self.redlock = Redlock([{"host": "localhost"}]) def test_lock(self): lock = self.redlock.lock("pants", 100) print(lock) self.assertEqual(lock.resource, "pants") self.redlock.unlock(lock) lock = self.redlock.lock("pants", 10) print(lock) self.redlock.unlock(lock)
def lightning_order_with_redlock() -> None: """ Lightning order with Redlock algorithm. :return: None """ r = redis.Redis() dlm = Redlock([{ 'host': 'localhost', 'port': 6379, 'db': 0 }, ]) # Stands for "distributed lock manager" lock = None try: # Try to acquire the lock lock = dlm.lock(LOCK_KEY, 30000) # If not acquiring the lock, block here # Business codes remaining = int(r.get('stock')) if remaining > 0: r.set('stock', str(remaining - 1)) print(f'Deducted stock, {remaining - 1} remaining') else: print('Failed to deduct stock') except MultipleRedlockException as e: print(e) finally: # Release the lock dlm.unlock(lock)
class RedisLock(object): def __init__(self, redis_host='127.0.0.1', redis_port=6379): self._lock_manager = Redlock([{ "host": redis_host, "port": redis_port, "db": 0, }]) @contextmanager def lock(self, lock_id, timeout=10000): ''' Context manager to acquire global distributed lock across all applications using the same redis. Uses implementation of redis distlock @see https://redis.io/topics/distlock @lock_id string uniq id of the lock @timeout integer timeout after which the lock will be released (in milliseconds) ''' lock = self._lock_manager.lock(lock_id, timeout) if lock is False: raise LockFailed() try: yield lock finally: self._lock_manager.unlock(lock)
def dlock(key, ttl, **kwargs): """ 分布式锁 :param key: 分布式锁ID :param ttl: 分布式锁生存时间 :param kwargs: 可选参数字典 :return: None """ resource_servers = [{ 'host': REDIS_HOST, 'port': REDIS_PORT, 'db': REDIS_DB, 'password': REDIS_PASSWORD }] dl = Redlock(resource_servers) # 获取锁 lock = dl.lock(key, ttl) # if ret is False: # detail = u'acquire lock[%s] error' % key # raise AcquireLockError(detail) yield lock # 释放锁 if isinstance(lock, Lock): dl.unlock(lock)
class RedisClient(object): def __init__(self, host, port=6379, db=0): self.redis = redis.StrictRedis(host=host, port=port, db=db) self.redlock = Redlock([self.redis]) def __getattr__(self, name): method = getattr(self.redis, name) return method def acquireDLock(self, resource, ttl=300): """ 获取分布式锁。 ttl: 锁的过期时间,单位是秒,默认300秒。 """ resource = 'lock_' + resource try: return self.redlock.lock(resource, ttl * 1000) except MultipleRedlockException as e: logging.error('Error acquiring dlock: %s', e) return False def releaseDLock(self, lock): """ 释放分布式锁 """ try: self.redlock.unlock(lock) return True except MultipleRedlockException as e: logging.error('Error releasing dlock: %s', e) return False
class TestRedlock(unittest.TestCase): def setUp(self): try: self.redlock = Redlock([{"host": "localhost"}]) self.dstlock = Redlock([{"host": "localhost", "port": 6379, "socket_timeout": 0.5}, {"host": "localhost", "port": 6380, "socket_timeout": 0.5}, {"host": "localhost", "port": 6381, "socket_timeout": 0.5}]) except Exception as e: pass ''' def test_distribute_locks(self): import time while True: if not self.dstlock.lock("fizz", 1000): print ("lock failed") else: print("lock success") time.sleep(1) ''' def test_lock(self): lock = self.redlock.lock("pants", 100) self.assertEqual(lock.resource, "pants") self.redlock.unlock(lock) lock = self.redlock.lock("pants", 10) self.redlock.unlock(lock) def test_blocked(self): lock = self.redlock.lock("pants", 1000) bad = self.redlock.lock("pants", 10) self.assertFalse(bad) self.redlock.unlock(lock) def test_bad_connection_info(self): with self.assertRaises(Warning): Redlock([{"cat": "hog"}]) def test_py3_compatible_encoding(self): lock = self.redlock.lock("pants", 1000) key = self.redlock.servers[0].get("pants") self.assertEqual(lock.key, key) def test_ttl_not_int_trigger_exception_value_error(self): with self.assertRaises(ValueError): self.redlock.lock("pants", 1000.0) def test_multiple_redlock_exception(self): ex1 = Exception("Redis connection error") ex2 = Exception("Redis command timed out") exc = MultipleRedlockException([ex1, ex2]) exc_str = str(exc) self.assertIn('connection error', exc_str) self.assertIn('command timed out', exc_str)
class RedisStore(AbstractStore): """Redis-based backend for deployments with replicas > 1""" def __init__(self, url: str): logger.info("Connecting to Redis on {}..".format(url)) self._redis = redis.StrictRedis.from_url(url) self._redlock = Redlock([url]) def set(self, key, value): self._redis.set(key, json.dumps(value, separators=(",", ":"))) def get(self, key): value = self._redis.get(key) if value: return json.loads(value.decode("utf-8")) def acquire_lock(self): return self._redlock.lock("update", 10000) def release_lock(self, lock): self._redlock.unlock(lock) def publish(self, event_type, event_data): self._redis.publish( "default", "{}:{}".format(event_type, json.dumps(event_data, separators=(",", ":"))), ) def listen(self): p = self._redis.pubsub() p.subscribe("default") for message in p.listen(): if message["type"] == "message": event_type, data = message["data"].decode("utf-8").split( ":", 1) yield (event_type, json.loads(data)) def create_screen_token(self): """Generate a new screen token and store it in Redis""" data = generate_token_data() token = data["token"] self._redis.set("screen-tokens:{}".format(token), json.dumps(data)) return token def redeem_screen_token(self, token: str, remote_addr: str): """Validate the given token and bind it to the IP""" redis_key = "screen-tokens:{}".format(token) data = self._redis.get(redis_key) if not data: raise ValueError("Invalid token") data = json.loads(data.decode("utf-8")) data = check_token(token, remote_addr, data) self._redis.set(redis_key, json.dumps(data))
def select_gpu(redis_conf=None, timeout=10000, shuffle=True): """Sets the CUDA_VISIBLE_DEVICES environment variable :param redis_conf: Redis configuration passed to redlock-py :param timeout: Timeout of the lock in milliseconds, default 10000 :param shuffle: Shuffles the available GPU list, default True """ if len(os.environ.get('CUDA_VISIBLE_DEVICES', 'unset')) == 0: # Environment variable empty return "", None gpu_status = str(subprocess.check_output(['nvidia-smi', 'pmon', '-c', '1'])) # Example of expected result from nvidia-smi: # # gpu pid type sm mem enc dec command # # Idx # C/G % % % % name # 0 25729 C 94 57 0 0 python # 1 - - - - - - - gpu_status = map(lambda x: x.split(), gpu_status.split("\n")[2:-1]) # Check if the GPU is not already used by the current process pid = os.getpid() gpu_pids = map(lambda x: x[:2], gpu_status) for gpu, p in gpu_pids: if p == '-': continue if pid == int(p): return int(gpu), None gpu_status = filter(lambda x: x[7] == '-', gpu_status) if shuffle: # Suffle GPUs list random.shuffle(gpu_status) if redis_conf is None: redis_conf = {'unix_socket_path': '/var/run/redis/redis-server.sock'} dlm = Redlock([redis_conf]) if len(gpu_status) > 0: for gpu_ in gpu_status: gpu = gpu_[0] gpu_lock = dlm.lock("{}:gpu{}".format(platform.node(), gpu), timeout) if gpu_lock != False: os.environ['CUDA_VISIBLE_DEVICES'] = gpu def unlock(): return dlm.unlock(gpu_lock) return int(gpu), unlock raise Exception("No GPU available!")
class RedisStore(AbstractStore): '''Redis-based backend for deployments with replicas > 1''' def __init__(self, url: str): logger.info('Connecting to Redis on {}..'.format(url)) self._redis = redis.StrictRedis.from_url(url) self._redlock = Redlock([url]) def set(self, key, value): self._redis.set(key, json.dumps(value, separators=(',', ':'))) def get(self, key): value = self._redis.get(key) if value: return json.loads(value.decode('utf-8')) def acquire_lock(self): return self._redlock.lock('update', 10000) def release_lock(self, lock): self._redlock.unlock(lock) def publish(self, event_type, event_data): self._redis.publish( 'default', '{}:{}'.format(event_type, json.dumps(event_data, separators=(',', ':')))) def listen(self): p = self._redis.pubsub() p.subscribe('default') for message in p.listen(): if message['type'] == 'message': event_type, data = message['data'].decode('utf-8').split( ':', 1) yield (event_type, json.loads(data)) def create_screen_token(self): '''Generate a new screen token and store it in Redis''' data = generate_token_data() token = data['token'] self._redis.set('screen-tokens:{}'.format(token), json.dumps(data)) return token def redeem_screen_token(self, token: str, remote_addr: str): '''Validate the given token and bind it to the IP''' redis_key = 'screen-tokens:{}'.format(token) data = self._redis.get(redis_key) if not data: raise ValueError('Invalid token') data = json.loads(data.decode('utf-8')) data = check_token(token, remote_addr, data) self._redis.set(redis_key, json.dumps(data))
def id(self): deviceIP = self.client.host self.session_id = self.get_redis_session(deviceIP) if not self.check_session_invalid(self.session_id, deviceIP): local_ip = self.get_local_ip() redisobj = RedisConfig() redis_server = redisobj.get_redis_servers(True) redis_conf_ls = [{ "host": s["host"], "port": s["port"], "db": s["dbno"] } for s in redis_server] lock_mgmt = Redlock(redis_conf_ls) device_lock = lock_mgmt.lock( local_ip + "_" + deviceIP + '_' + self.username + '_device_lock_calabash', 30 * 1000) tmp_count = 0 while isinstance(device_lock, bool) and not device_lock and tmp_count < 1000: tmp_count += 1 time.sleep(0.5) self.session_id = self.get_redis_session(deviceIP) if self.check_session_invalid(self.session_id, deviceIP): return self.session_id device_lock = lock_mgmt.lock( local_ip + "_" + deviceIP + '_' + self.username + '_device_lock_calabash', 30 * 1000) self.session_id = self.get_redis_session(deviceIP) if self.check_session_invalid(self.session_id, deviceIP): return self.session_id self.authenticate(self.username, self.password) #set sessionid to redis self.set_session_to_redis(self.session_id, deviceIP) lock_mgmt.unlock(device_lock) return self.session_id
class Bloom: def __init__(self, key, host='127.0.0.1', port=6379, db=0): self.key = key # redis-bitmap的key self.redis_cli = redis.StrictRedis(host=host, port=port, db=db, charset='utf-8') self.red_lock = Redlock([ { "host": host, "port": port, "db": db }, ]) def add(self, value): value_unique_lock = self.red_lock.lock(value, 60 * 1000) if value_unique_lock is False: return False point_list = self.get_positions(value) for b in point_list: exists = self.redis_cli.getbit(self.key, b) if not exists: break else: return False for b in point_list: self.redis_cli.setbit(self.key, b, 1) return True def contains(self, value): point_list = self.get_positions(value) for b in point_list: exists = self.redis_cli.getbit(self.key, b) if not exists: return False return True @staticmethod def get_positions(value): point0 = mmh3.hash(value, 40, signed=False) point1 = mmh3.hash(value, 41, signed=False) point2 = mmh3.hash(value, 42, signed=False) point3 = mmh3.hash(value, 43, signed=False) point4 = mmh3.hash(value, 44, signed=False) point5 = mmh3.hash(value, 45, signed=False) point6 = mmh3.hash(value, 46, signed=False) point7 = mmh3.hash(value, 47, signed=False) return [point0, point1, point2, point3, point4, point5, point6, point7]
class RedisStore(AbstractStore): '''Redis-based backend for deployments with replicas > 1''' def __init__(self, url: str): logger.info('Connecting to Redis on {}..'.format(url)) self._redis = redis.StrictRedis.from_url(url) self._redlock = Redlock([url]) def set(self, key, value): self._redis.set(key, json.dumps(value, separators=(',', ':'))) def get(self, key): value = self._redis.get(key) if value: return json.loads(value.decode('utf-8')) def acquire_lock(self): return self._redlock.lock('update', 10000) def release_lock(self, lock): self._redlock.unlock(lock) def publish(self, event_type, event_data): self._redis.publish('default', '{}:{}'.format(event_type, json.dumps(event_data, separators=(',', ':')))) def listen(self): p = self._redis.pubsub() p.subscribe('default') for message in p.listen(): if message['type'] == 'message': event_type, data = message['data'].decode('utf-8').split(':', 1) yield (event_type, json.loads(data)) def create_screen_token(self): '''Generate a new screen token and store it in Redis''' data = generate_token_data() token = data['token'] self._redis.set('screen-tokens:{}'.format(token), json.dumps(data)) return token def redeem_screen_token(self, token: str, remote_addr: str): '''Validate the given token and bind it to the IP''' redis_key = 'screen-tokens:{}'.format(token) data = self._redis.get(redis_key) if not data: raise ValueError('Invalid token') data = json.loads(data.decode('utf-8')) data = check_token(token, remote_addr, data) self._redis.set(redis_key, json.dumps(data))
class TestRedlock(unittest.TestCase): def setUp(self): self.redlock = Redlock([{"host": "localhost"}]) def test_lock(self): lock_name = "test_lock" lock = self.redlock.lock(lock_name, 100, TEST_LOCK_TTL_MILLIS) self.assertEqual(lock.resource, lock_name) self.redlock.unlock(lock) lock = self.redlock.lock(lock_name, 10, TEST_LOCK_TTL_MILLIS) self.redlock.unlock(lock) def test_blocked(self): lock_name = "test_blocked" lock = self.redlock.lock(lock_name, 1000, TEST_LOCK_TTL_MILLIS) with self.assertRaises(MultipleRedlockException): self.redlock.lock(lock_name, 10, TEST_LOCK_TTL_MILLIS) self.redlock.unlock(lock) def test_bad_connection_info(self): with self.assertRaises(Warning): Redlock([{"cat": "hog"}]) def test_py3_compatible_encoding(self): lock_name = "test_py3_compatible_encoding" lock = self.redlock.lock(lock_name, 1000, TEST_LOCK_TTL_MILLIS) key = self.redlock.servers[0].get(lock_name) self.assertEqual(lock.key, int(key)) def test_ttl_not_int_trigger_exception_value_error(self): lock_name = "ttl_not_int" with self.assertRaises(ValueError): self.redlock.lock(lock_name, 1000.0, 1000.0) def test_multiple_redlock_exception(self): ex1 = Exception("Redis connection error") ex2 = Exception("Redis command timed out") exc = MultipleRedlockException([ex1, ex2]) exc_str = str(exc) self.assertIn('connection error', exc_str) self.assertIn('command timed out', exc_str)
class GPUManager: """Class for allocating GPUs.""" def __init__(self, verbose: bool = False): self.lock_manager = Redlock([ { "host": "localhost", "port": 6379, "db": 0 }, ]) self.verbose = verbose def get_free_gpu(self): """ If some GPUs are available, try reserving one by checking out an exclusive redis lock. If none available or can't get lock, sleep and check again. """ while True: gpu_ind = self._get_free_gpu() if gpu_ind is not None: return gpu_ind if self.verbose: print(f'pid {os.getpid()} sleeping') time.sleep(GPU_LOCK_TIMEOUT / 1000) def _get_free_gpu(self): try: available_gpu_inds = [ gpu.index for gpu in gpustat.GPUStatCollection.new_query() if gpu.memory_used < 0.5 * gpu.memory_total ] except Exception: return [0] # Return dummy GPU index if no CUDA GPUs are installed if available_gpu_inds: gpu_ind = np.random.choice(available_gpu_inds) if self.verbose: print(f'pid {os.getpid()} picking gpu {gpu_ind}') if self.lock_manager.lock(f'gpu_{gpu_ind}', GPU_LOCK_TIMEOUT): return int(gpu_ind) if self.verbose: print(f'pid {os.getpid()} couldnt get lock') return None
class SlaveSpider(Spider): name = 'slave_spider' default_retry_count = 3 default_retry_delay = 1 def __init__(self): self.dlm = Redlock([{"host": "42.96.132.158", "port": 6379, "db": 0, 'password': "******"}, ]) self.list = [] self.driver = PhantomJS() def start_requests(self): self.load_payload() while True: my_lock = self.dlm.lock('URLSLOCK', 1000) while isinstance(my_lock, bool): pass len = self.dlm.servers[0].llen('urls') while len > 0: url = self.dlm.servers[0].lpop('urls') self.dlm.unlock(my_lock) urls = self.get_urls(url=url) for u in urls: yield Request(url=u, callback=self.parse) log.msg('the slave spider is over!', log.WARNING) def parse(self, response): log.msg(response.url+' is over', log.INFO) def get_urls(self, url): for l in self.list: temp = url + l yield temp def load_payload(self): f = open('Xspider/payload.txt', 'r') for line in f.readlines(): line = line.strip('\n') self.list.append(line) f.close()
class RedisLocker(ILocker): def __init__(self, redis_con): self.dlm = Redlock([redis_con, ], retry_count=1) self._locks = {} def _lock(self, key, timeout): """ Trying to acquire lock If not success then sleep (timeout // 10) milliseconds and try again """ lock_key = f'lock/{key}' sleep_time = timeout / (10 * 1000) while True: lock = self.dlm.lock(lock_key, timeout) if lock: break time.sleep(sleep_time) self._locks[lock_key] = lock def _unlock(self, key): lock_key = f'lock/{key}' lock = self._locks[lock_key] self.dlm.unlock(lock)
class TestRedlock(unittest.TestCase): def setUp(self): self.redlock = Redlock([{"host": "localhost"}]) def test_lock(self): lock = self.redlock.lock("pants", 100) self.assertEqual(lock.resource, "pants") self.redlock.unlock(lock) lock = self.redlock.lock("pants", 10) self.redlock.unlock(lock) def test_blocked(self): lock = self.redlock.lock("pants", 1000) bad = self.redlock.lock("pants", 10) self.assertFalse(bad) self.redlock.unlock(lock) def test_bad_connection_info(self): with self.assertRaises(Warning): Redlock([{"cat": "hog"}]) def test_py3_compatible_encoding(self): lock = self.redlock.lock("pants", 1000) key = self.redlock.servers[0].get("pants") self.assertEquals(lock.key, key) def test_ttl_not_int_trigger_exception_value_error(self): with self.assertRaises(ValueError): self.redlock.lock("pants", 1000.0) def test_multiple_redlock_exception(self): ex1 = Exception("Redis connection error") ex2 = Exception("Redis command timed out") exc = MultipleRedlockException([ex1, ex2]) exc_str = str(exc) self.assertIn('connection error', exc_str) self.assertIn('command timed out', exc_str)
class NLReceiver: def __init__(self, nl_rest_url, data_key, user="******"): self.user = user self._dlm = Redlock([ { "host": "redis_trade", "port": 6379, "db": 0 }, ]) self.error = None self.nl_rest_url = nl_rest_url self.data_key = data_key self.token = self._get_token self.data = self._get_data def _pause(self): while True: my_lock = self._dlm.lock("pause_for_API", 200) if my_lock: break print("PAUSE_FOR_API", self.user) def _connector(self, nl_rest_url, data_key, params): data_response = None data_result = {} data_result_error = None count = 0 def _msg_to_logger(msg, data_response_status): return "USER: [{}] PAUSE_FOR_API - {} , data_response_status[\"code\"] - {}," \ "data_response_status - {}, - nl_rest_url {}, data_key - {}, params - {}"\ .format(self.user, msg, data_response_status["code"], data_response_status, nl_rest_url, data_key, params) try: pause = False while count != 40: count += 1 self._pause() r = s.get(nl_rest_url, params=params) data_response = json.loads(r.text[6:]) # JSON Hijacking data_response_status = data_response[data_key]["status"] if data_response_status["code"] == "200": if pause: logger_app.error( _msg_to_logger("PAUSE GOOD", data_response_status)) data_result = data_response[data_key]["data"] break elif data_response_status["code"] == "4": if pause: logger_app.error( _msg_to_logger("PAUSE GOOD", data_response_status)) break pause = True logger_app.error( _msg_to_logger("PROBLEM", data_response_status)) params = self._get_token else: raise Exception(str(data_response)) except requests.exceptions.RequestException as e: msg = "USER: [{}] RequestException {}".format(self.user, str(e)) data_result_error = msg logger_app.error(msg) except Exception as e: msg = "USER: [{}] Exception - {} - {} - data_response: {}".format( self.user, e.__class__, str(e), data_response) data_result_error = msg logger_app.error(msg) # else: # # logger_dealer_nl.info("USER: [{}] {} {}".format(self.user, data_key, str(data_response_status))) # if data_response_status["code"] == "200": # data_response = data_response[data_key]["data"] finally: return {"data": data_result, "error": data_result_error} @property def _get_token(self): data_key = current_app.config["NL_AUTH"]["DATA_KEY"] nl_url_auth = current_app.config["NL_AUTH"]["URL"] params = { "username": current_app.config["NL_USERNAME"], "password": current_app.config["NL_PASSWORD"] } response = self._connector(nl_url_auth, data_key, params) data = response.get("data") token = data.get("token") if token: return {"oauth_token": token} if response.get("error"): self.error = response.get("error") @property def _get_data(self): if self.token: response = self._connector(self.nl_rest_url, self.data_key, self.token) self.error = response.get("error") return response.get("data") else: return {} def __str__(self): return "error: {}, nl_rest_url: {}, data_key: {}, token: {}, data: {}, user: {}" \ .format(self.error, self.nl_rest_url, self.data_key, self.token, self.data, self.user,)
# Get the resources previously to its modification print(f"{bcolors.OKCYAN}Retrieving resource '" + RESOURCE_1_NAME + f"'in order to compare it after the update..." f"{bcolors.ENDC}") prev_resource_1 = redis.get(RESOURCE_1_NAME) print(f"{bcolors.OKCYAN}Retrieving resource '" + RESOURCE_2_NAME + f"'in order to compare it after the update..." f"{bcolors.ENDC}") prev_resource_2 = redis.get(RESOURCE_2_NAME) # Get a lock on the plane number 0 around 10 seconds print(f"{bcolors.OKGREEN}Client 1 - Lock 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Requesting lock on resource '" + RESOURCE_1_NAME + f"'...{bcolors.ENDC}") plane_0_lock_1 = dlm.lock(RESOURCE_1_ID, 10000) # Get another lock on the plane number 1 around 5 seconds print(f"{bcolors.OKGREEN}Client 1 - Lock 2:{bcolors.ENDC} " f"{bcolors.OKCYAN}Requesting lock on resource '" + RESOURCE_2_NAME + f"'...{bcolors.ENDC}") plane_1_lock_1 = dlm.lock(RESOURCE_2_ID, 5000) # Check if the lock has been acquired if plane_0_lock_1: print(f"{bcolors.OKGREEN}Client 1 - Lock 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Lock on resource '" + RESOURCE_1_NAME + f"' acquired successfully. {bcolors.ENDC}") lock_key = plane_0_lock_1.key.decode("utf-8") print(
class ProbitScheduler(Scheduler): def __init__(self, redis_connection=None, locker=None, *args, **kwargs): self.__redis_connection = redis_connection if self.__redis_connection is None: self.__redis_connection = StrictRedis.from_url( current_app.conf.CELERY_REDIS_SCHEDULER_URL) self._schedule = EntryProxy(self.__redis_connection) self._locker = locker if self._locker is None: self._locker = Redlock( [current_app.conf.CELERY_REDIS_SCHEDULER_URL]) super(ProbitScheduler, self).__init__(*args, **kwargs) def setup_schedule(self): self.install_default_entries(self._schedule) self._merge(self.app.conf.CELERYBEAT_SCHEDULE) def get_schedule(self): return self._schedule schedule = property( get_schedule) # This isn't inherited anymore? Do we want to do this? def sync(self): # Reload the schedule from the collection self._schedule = EntryProxy(self.__redis_connection) # I'm not sure what reserve() is intended to do, but it does not do what we # need it to do, so we define a _lock() method as well. def maybe_due(self, entry, publisher=None): is_due, next_time_to_run = entry.is_due() if not is_due: return next_time_to_run lock = self._lock(entry.name) if not lock: return next_time_to_run try: # Now that we have the lock, double-check the timestamps on the # entry before executing it. entry = self._schedule.sync(entry.name) if entry is None: return next_time_to_run is_due, next_time_to_run = entry.is_due() if not is_due: return next_time_to_run return Scheduler.maybe_due(self, entry, publisher) finally: self._unlock(lock) def _lock(self, name): return self._locker.lock(name, 1000) def _unlock(self, lock): self._locker.unlock(lock) def _merge(self, schedule): """schedule_keys = self.__redis_connection.hgetall(ENTRY_LIST_KEY).keys() if len(schedule_keys) > 0: self.__redis_connection.hdel(ENTRY_LIST_KEY, *schedule_keys)""" for name, entry_dict in list(schedule.items()): entry = ScheduleEntry(name, **entry_dict) if name not in self._schedule: self._schedule[name] = entry else: # _lock() the existing entry so that these values aren't changed # while we're merging them. lock = self._lock(name) if lock: try: existing = self._schedule.sync(name) if existing: entry.last_run_at = existing.last_run_at entry.total_run_count = existing.total_run_count self._schedule[name] = entry finally: self._unlock(lock)
# Get the resource previously to its modification print(f"{bcolors.OKCYAN}Retrieving resource '" + RESOURCE_1_NAME + f"'in order to compare it after the update..." f"{bcolors.ENDC}") prev_resource_1 = redis.get(RESOURCE_1_NAME) print(f"{bcolors.OKCYAN}Retrieving resource '" + RESOURCE_2_NAME + f"'in order to compare it after the update..." f"{bcolors.ENDC}") prev_resource_2 = redis.get(RESOURCE_2_NAME) # Get a lock on the plane number 1 around 10 seconds print(f"{bcolors.OKGREEN}Client 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Requesting lock on resource '" + RESOURCE_1_NAME + f"'...{bcolors.ENDC}") plane_0_lock_1 = dlm.lock(RESOURCE_1_ID, 10000) # Start the second client process client_2 = threading.Thread(target=thread_function) # Start the second client with the thread print(f"{bcolors.OKGREEN}Client 2: Starting...{bcolors.ENDC} ") client_2.start() # Check if the lock has been acquired if plane_0_lock_1: print(f"{bcolors.OKGREEN}Client 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Lock on resource '" + RESOURCE_1_NAME + f"' acquired successfully. {bcolors.ENDC}") lock_key = plane_0_lock_1.key.decode("utf-8")
# host = '42.96.132.158' # psd = 'wjdh84928399' # r = redis.Redis(host=host, password=psd, db=0, port=6379, socket_timeout=10) # url = r.lpop('urls') # f = open('Xspider/payload.txt', 'r') # list = [] # for line in f.readlines(): # line = line.strip('\n') # url = url + line # list.append(url) # print list from redlock import Redlock from redlock import MultipleRedlockException import time dlm = Redlock([{"host": "42.96.132.158", "port": 6379, "db": 0, 'password': "******"}, ]) try: dlm.servers[0].flushall() my_lock = dlm.lock('LOCK', 1000) while True: if isinstance(my_lock, bool): print 'wait' time.sleep(0.5) else: print 'dosomething' dlm.servers[0].lpush('urls', time.time()) dlm.unlock(my_lock) time.sleep(0.5) except MultipleRedlockException, e: raise e
plugins = ['time', 'choice', 'var'] analyzers = [] for plugin in plugins: mod = importlib.import_module('plugins.' + plugin) Analyzer = getattr(mod, 'Analyzer') analyzers.append(Analyzer()) gameplays = db.gameplays.find({'modified': True}) for gameplay in gameplays: try: gameplay_file = '%s' % gameplay['_id'] lock = False while not lock: lock = dlm.lock(gameplay_file, 5000) try: shutil.move(traces_path + gameplay_file, cache_path + gameplay_file) except FileNotFoundError as e: print(e) continue finally: dlm.unlock(lock) db.gameplays.update_one({'_id': gameplay['_id'], 'lastAccessed': gameplay['lastAccessed']}, {'$set': {'modified': False}}) with open(cache_path + gameplay_file, 'rb') as src: with open(storage_path + gameplay_file, 'ab') as dst: dst.write(src.read()) with open(cache_path + gameplay_file, 'r') as f:
class ConsumerQueue(object): def __init__(self, redis_host=None, redis_port=None, redis_db=None, producer_queue_id=None, consumer_queue_id=None): assert redis_host assert redis_port assert redis_db assert producer_queue_id assert consumer_queue_id pool = redis.ConnectionPool(host=redis_host, port=redis_port, db=redis_db) self._broker = redis.Redis(connection_pool=pool) self._consumer_queue_id = producer_queue_id self._consumer_queue_id = consumer_queue_id self._producer_queue_id = producer_queue_id self._lock_manager = Redlock([{ 'host': redis_host, 'port': redis_port, 'db': redis_db}, ]) def dequeue_task(self): """ Pulls the descriptor for a task that is ready for processing returns a descriptor or None if no tasks are present """ # do we still have work items in this worker's queue descriptor = self._broker.lindex(self._consumer_queue_id, 0) if not descriptor: # check the global work queue for work descriptor = self._broker.rpoplpush(self._producer_queue_id, self._consumer_queue_id) return descriptor def retire_task(self, descriptor=None): """ Mark a task as having been successfully dispatched """ assert descriptor popped_id = self._broker.lpop(self._consumer_queue_id) assert popped_id == descriptor return popped_id def lock_work_item(self, descriptor=None, lock_timeout=None): """ Locks a task to prevent other consumers from double processing """ assert descriptor assert lock_timeout lock = self._lock_manager.lock(descriptor, lock_timeout) return lock def unlock_task(self, lock=None): assert lock self._lock_manager.unlock(lock)
class ProbitScheduler(Scheduler): def __init__(self, redis_connection=None, locker=None, *args, **kwargs): self.__redis_connection = redis_connection if self.__redis_connection is None: self.__redis_connection = StrictRedis.from_url(current_app.conf.CELERY_REDIS_SCHEDULER_URL) self._schedule = EntryProxy(self.__redis_connection) self._locker = locker if self._locker is None: self._locker = Redlock([current_app.conf.CELERY_REDIS_SCHEDULER_URL]) super(ProbitScheduler, self).__init__(*args, **kwargs) def setup_schedule(self): self.install_default_entries(self._schedule) self._merge(self.app.conf.CELERYBEAT_SCHEDULE) def get_schedule(self): return self._schedule schedule = property(get_schedule) # This isn't inherited anymore? Do we want to do this? def sync(self): # Reload the schedule from the collection self._schedule = EntryProxy(self.__redis_connection) # I'm not sure what reserve() is intended to do, but it does not do what we # need it to do, so we define a _lock() method as well. def maybe_due(self, entry, publisher=None): is_due, next_time_to_run = entry.is_due() if not is_due: return next_time_to_run lock = self._lock(entry.name) if not lock: return next_time_to_run try: # Now that we have the lock, double-check the timestamps on the # entry before executing it. entry = self._schedule.sync(entry.name) if entry is None: return next_time_to_run is_due, next_time_to_run = entry.is_due() if not is_due: return next_time_to_run return Scheduler.maybe_due(self, entry, publisher) finally: self._unlock(lock) def _lock(self, name): return self._locker.lock(name, 1000) def _unlock(self, lock): self._locker.unlock(lock) def _merge(self, schedule): """schedule_keys = self.__redis_connection.hgetall(ENTRY_LIST_KEY).keys() if len(schedule_keys) > 0: self.__redis_connection.hdel(ENTRY_LIST_KEY, *schedule_keys)""" for name, entry_dict in schedule.items(): entry = ScheduleEntry(name, **entry_dict) if name not in self._schedule: self._schedule[name] = entry else: # _lock() the existing entry so that these values aren't changed # while we're merging them. lock = self._lock(name) if lock: try: existing = self._schedule.sync(name) if existing: entry.last_run_at = existing.last_run_at entry.total_run_count = existing.total_run_count self._schedule[name] = entry finally: self._unlock(lock)
class CacheManager(object): def __init__(self, redis_engine, lock_ttl=1000): self.redis = redis_engine self.redlock = Redlock([redis_engine]) self.lock_ttl = lock_ttl @staticmethod def get_lock_key(key): return "%s:lock" % key def lock(self, key, nowait=False): lock_key = self.get_lock_key(key) l = False while not l: l = self.redlock.lock(lock_key, self.lock_ttl) if l or nowait: return l time.sleep(0.3) def unlock(self, lock): self.redlock.unlock(lock) def get_item(self, key, recover=None, ttl=None, arguments=([], {}), **kw): item_key = self.redis.get(key) value = self.redis.get(item_key) if item_key else None # return None if cannot recover the value if value is None and recover is None: return None # Trying to recover the value if value is None: return self.set_item(key, lambda: recover(*arguments[0], **arguments[1]), ttl=ttl) else: return deserialize(value) def set_item(self, key, value, ttl=None): old_item_key = self.redis.get(key) if old_item_key is None: # this is a new item # locking it to prevent concurrency violation lock = self.lock(key, nowait=True) if not lock: # it seems this item is loading in another thread # so, waiting for that: # wait & make sure the object is reloaded, the release the lock self.unlock(self.lock(key)) return self.get_item(key) # check if item loaded v = self.get_item(key, None) if v: self.unlock(lock) return v else: lock = None try: guid = uuid.uuid1() item_key = "%s:%s" % (key, guid) value = value() if callable(value) else value self.redis.set(item_key, serialize(value), ex=ttl) self.redis.set(key, item_key, ex=ttl) if old_item_key: self.redis.delete(old_item_key) return value finally: if lock: self.unlock(lock) def hget_item(self, key, dict_key, recover=None, ttl=None, arguments=([], {}), **kw): item_key = self.redis.get(key) value = self.redis.hget(item_key, dict_key) if item_key else None # return None if cannot recover the value if value is None and recover is None: return None # Trying to recover the value if value is None: return self.hset_item(key, dict_key, lambda: recover(*arguments[0], **arguments[1]), ttl=ttl) else: return deserialize(value) def hset_item(self, key, dict_key, value, ttl=None): old_item_key = self.redis.get(key) if old_item_key is None: # this is a new item # locking it to prevent concurrency violation lock = self.lock(key, nowait=True) if not lock: # it seems this item is loading in another thread # so, waiting for that: # wait & make sure the object is reloaded, the release the lock self.unlock(self.lock(key)) return self.hget_item(key, dict_key) # check if item loaded v = self.hget_item(key, dict_key, None) if v: self.unlock(lock) return v else: lock = None try: guid = uuid.uuid1() item_key = "%s:%s" % (key, guid) value = value() if callable(value) else value self.redis.hset(item_key, dict_key, serialize(value)) self.redis.expire(item_key, time=ttl) self.redis.set(key, item_key, ex=ttl) if old_item_key: self.redis.delete(old_item_key) return value finally: if lock: self.unlock(lock) def del_item(self, *keys): item_keys_to_remove = [] for key in keys: item_key = self.redis.get(key) if item_key is not None: item_keys_to_remove.append(item_key) self.redis.delete(*item_keys_to_remove) self.redis.delete(*keys) def get_list(self, key, recover=None, ttl=None, arguments=([], {}), key_extractor=None): # First, get the keys list value = self.redis.get(key) if value is None and recover is None: return None, None # Trying to recover the value if value is None: return self.set_list( key, lambda: recover(*arguments[0], **arguments[1]), ttl=ttl, key_extractor=key_extractor ) else: metadata, item_keys = deserialize(value) return metadata, [self.get_item(item_key) for item_key in item_keys] def set_list(self, key, value, ttl=None, key_extractor=None): # locking it to prevent concurrency violation lock = self.lock(key) # check if item loaded v = self.get_list(key) if v[1]: self.unlock(lock) return v # packed: (value_metadata, value) if key_extractor is None: def key_extractor(x): return x["id"] try: value_metadata, value = value() if callable(value) else value item_keys = [] assert isinstance(value, collections.Iterable), "Value must be iterable" for item_value in value: item_key = create_cache_key(key.split(":")[0], key_extractor(item_value)) item_keys.append(item_key) self.set_item(item_key, item_value, ttl=ttl) self.redis.set(key, serialize((value_metadata, item_keys)), ex=ttl) return value_metadata, value finally: if lock: self.unlock(lock) def invalidate_item(self, key): item_key = self.redis.get(key) self.redis.delete(key) if item_key: self.redis.delete(item_key) def invalidate_list(self, key): self.redis.delete(key) def decorate(self, namespace, list_=False, ttl=None, key_extractor=None): def decorator(func): def wrapper(*args, **kwargs): cache_key = create_cache_key(namespace, kwargs) cache_params = dict(ttl=ttl, recover=func, arguments=(args, kwargs)) if list_: cache_method = self.get_list cache_params["key_extractor"] = key_extractor else: cache_method = self.get_item return cache_method(cache_key, **cache_params) return wrapper return decorator
plugins = ['time', 'choice', 'var'] analyzers = [] for plugin in plugins: mod = importlib.import_module('plugins.' + plugin) Analyzer = getattr(mod, 'Analyzer') analyzers.append(Analyzer()) gameplays = db.gameplays.find({'modified': True}) for gameplay in gameplays: try: gameplay_file = '%s' % gameplay['_id'] lock = False while not lock: lock = dlm.lock(gameplay_file, 5000) try: shutil.move(traces_path + gameplay_file, cache_path + gameplay_file) except FileNotFoundError as e: print(e) continue finally: dlm.unlock(lock) db.gameplays.update_one( { '_id': gameplay['_id'], 'lastAccessed': gameplay['lastAccessed'] }, {'$set': { 'modified': False
print(f"{bcolors.OKBLUE}## EXECUTING TEST 6 ##{bcolors.ENDC}") print( f"{bcolors.OKBLUE} Several clients, Several locks, one resource, client blocked {bcolors.ENDC}" ) # Get the resource previously to its modification print( f"{bcolors.OKCYAN}Retrieving resource in order to compare it after the update...{bcolors.ENDC}" ) prev_resource = redis.get(RESOURCE_NAME) # Get a lock on the plane number 1 around 10 seconds print(f"{bcolors.OKGREEN}Client 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Requesting lock on resource '" + RESOURCE_NAME + f"'...{bcolors.ENDC}") plane_lock = dlm.lock(RESOURCE_ID, 10000) # Start the second client process client_2 = threading.Thread(target=thread_function) # Start the second client with the thread print(f"{bcolors.OKGREEN}Client 2: Starting...{bcolors.ENDC} ") client_2.start() # Sleep current thread for 15 seconds in order to simulate it is blocked and check the behaviour print(f"{bcolors.OKGREEN}Client 1: Blocked for 15 seconds...{bcolors.ENDC} ") time.sleep(15) # Start the second client process, again, to try again client_2 = threading.Thread(target=thread_function)
class RecordLockManager(object): def __init__(self, conf, prefix=None, ttl_sec=60, backlog_path=None): self.backlog_lock = Lock() self.conf = conf self.redlock = Redlock([{"host": conf['backend']['host'], "port": int(conf['backend']['port']), "db": 13, "password": conf['backend']['password']}, ], retry_count=100000, retry_delay=0.01) if prefix is not None: self.prefix = prefix + "-" self.ttl = ttl_sec * 1000 self.backlog_filename_pat = re.compile(self.prefix + '([\w\d\.\-\\/]+?)\-backlog\.([\w\d\-]+)') self.backlog = {} self.backlog_path = backlog_path # Should be abs path to the dir containing backlog pickles self.is_dumping = False if not os.path.exists(self.backlog_path): os.makedirs(self.backlog_path) self.reload() def parse_backlog_filename(self, name): match = self.backlog_filename_pat.search(name) if match is not None: return match.group(1), match.group(2) else: return None, None def reload(self): for i in os.listdir(self.backlog_path): if os.path.isfile(os.path.join(self.backlog_path, i)): name, postfix = self.parse_backlog_filename(i) if name is not None: with open(self.make_real_path(i), 'rb') as ifile: a = pickle.load(ifile) (data, saved_context) = a task = self.make_task(name, data, postfix=postfix) task.saved_context = saved_context task.is_backlog = True self.backlog[task.pickle_name] = task def make_real_path(self, pickle_name): if self.backlog_path is not None: return os.path.join(self.backlog_path, pickle_name) else: return pickle_name # TODO: implement exponential back off def choose_backlog_and_lock_it(self, ttl_sec=None, retry=5): with self.backlog_lock: count = 0 while count < retry: if len(self.backlog) < 1: return None backlog_task = self.backlog[random.choice(self.backlog.keys())] if backlog_task.lock is not None: # this task is already locked, maybe by other thread continue count+=1 self.try_lock(backlog_task, ttl_sec) if not backlog_task.lock: return None else: return backlog_task return None def try_lock(self, task, ttl_sec=None, block=False): if ttl_sec is not None: ttl = ttl_sec * 1000 else: ttl = self.ttl lock = self.redlock.lock(task.lock_name, ttl) while block and not lock: lock = self.redlock.lock(task.lock_name, ttl) if not lock: task.lock = None else: task.lock = lock def store_backlog(self, task, saved_context): print "Store backlog " + task.lock_name self.is_dumping = True task.saved_context = saved_context with open(self.make_real_path(task.pickle_name), "wb") as ofile: pickle.dump((task.data, saved_context), ofile) self.is_dumping = False with self.backlog_lock: # Must ensure task holds no lock before store it as backlog # otherwise, it may never scheduled again if task.lock is not None: self.unlock(task.lock) task.lock = None task.is_backlog = True self.backlog[task.pickle_name] = task def make_task(self, name, data, postfix=None): return Task(name, self.prefix, data, is_backlog=False, postfix=postfix) def unlock(self, lock): self.redlock.unlock(lock) def commit_task(self, task): if task.is_backlog: with self.backlog_lock: if task.pickle_name in self.backlog: del self.backlog[task.pickle_name] real_path = self.make_real_path(task.pickle_name) if os.path.exists(real_path): os.remove(real_path) if task.lock is not None: self.unlock(task.lock) task.lock = None
# line = line.strip('\n') # url = url + line # list.append(url) # print list from redlock import Redlock from redlock import MultipleRedlockException import time dlm = Redlock([ { "host": "42.96.132.158", "port": 6379, "db": 0, 'password': "******" }, ]) try: dlm.servers[0].flushall() my_lock = dlm.lock('LOCK', 1000) while True: if isinstance(my_lock, bool): print 'wait' time.sleep(0.5) else: print 'dosomething' dlm.servers[0].lpush('urls', time.time()) dlm.unlock(my_lock) time.sleep(0.5) except MultipleRedlockException, e: raise e
if len(sys.argv) < 4: print('Invoke the script in the form "python kafka-consumer.py kafka_connection_string topic traces_path"') exit() kafka_connection = sys.argv[1] topic = sys.argv[2] traces_path = sys.argv[3] if not traces_path[-1] == '/': traces_path += "/" # To consume messages consumer = KafkaConsumer(topic, group_id='traces_cache', bootstrap_servers=[kafka_connection]) consumer.commit() dlm = Redlock([{"host": "localhost", "port": 6379, "db": 0}, ]) for message in consumer: key = message.key.decode('utf-8') lock = False while not lock: lock = dlm.lock(key, 5000) with open(traces_path + key, 'ab') as f: f.write(message.value) dlm.unlock(lock) consumer.task_done(message) consumer.commit()
traces_path = sys.argv[3] if not traces_path[-1] == '/': traces_path += "/" # To consume messages consumer = KafkaConsumer(topic, group_id='traces_cache', bootstrap_servers=[kafka_connection]) consumer.commit() dlm = Redlock([ { "host": "localhost", "port": 6379, "db": 0 }, ]) for message in consumer: key = message.key.decode('utf-8') lock = False while not lock: lock = dlm.lock(key, 5000) with open(traces_path + key, 'ab') as f: f.write(message.value) dlm.unlock(lock) consumer.task_done(message) consumer.commit()
print(f"{bcolors.OKBLUE}## EXECUTING TEST 2 ##{bcolors.ENDC}") print( f"{bcolors.OKBLUE} One client, Several locks, One resource {bcolors.ENDC}") # Get the resource previously to its modification print( f"{bcolors.OKCYAN}Retrieving resource in order to compare it after the update...{bcolors.ENDC}" ) prev_resource = redis.get(RESOURCE_NAME) # Get a lock on the plane number 1 around 10 seconds print(f"{bcolors.OKGREEN}Client 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Requesting first lock on resource '" + RESOURCE_NAME + f"'...{bcolors.ENDC}") plane_1_lock_1 = dlm.lock(RESOURCE_ID, 10000) # Get another lock on the plane number 1 around 5 seconds print(f"{bcolors.OKGREEN}Client 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Requesting second lock on resource '" + RESOURCE_NAME + f"'...{bcolors.ENDC}") plane_1_lock_2 = dlm.lock(RESOURCE_ID, 5000) # Check if the lock has been acquired if plane_1_lock_1 and not plane_1_lock_2: print(f"{bcolors.OKGREEN}Client 1 - Lock 1:{bcolors.ENDC} " f"{bcolors.OKCYAN}Lock on resource '" + RESOURCE_NAME + f"' acquired successfully. {bcolors.ENDC}") print(f"{bcolors.OKGREEN}Client 1 - Lock 2:{bcolors.ENDC} " f"{bcolors.FAIL}Lock on resource '" + RESOURCE_NAME + f"' not acquired. {bcolors.ENDC}")