def cb_no_overlap():
with Lock(conn, "foobar"):
time.sleep(0.001)
def test_get_owner_id(conn):
lock = Lock(conn, "foobar-tok")
lock.acquire()
assert lock.get_owner_id() == lock.id
lock.release()
def test_bogus_release(conn):
lock = Lock(conn, "foobar-tok")
pytest.raises(NotAcquired, lock.release)
lock.acquire()
lock2 = Lock(conn, "foobar-tok", id=lock.id)
lock2.release()
def test_expire_less_than_timeout(conn):
lock = Lock(conn, "foobar", expire=1)
with pytest.raises(TimeoutTooLarge):
lock.acquire(blocking=True, timeout=2)
def test_plain(conn):
with Lock(conn, "foobar"):
time.sleep(0.01)
def test_timeout(conn):
with Lock(conn, "foobar"):
lock = Lock(conn, "foobar")
assert lock.acquire(timeout=1) is False
def test_timeout_expire_with_renewal(conn):
with Lock(conn, "foobar", expire=1, auto_renewal=True):
lock = Lock(conn, "foobar")
assert lock.acquire(timeout=2) is False
def test_signal_cleanup_on_reset_all(conn):
"""After resetting all locks, no signal keys should not remain."""
lock = Lock(conn, 'foo')
lock.acquire()
reset_all(conn)
assert conn.llen('lock-signal:foo') == 0
def test_double_acquire(redis_server):
lock = Lock(StrictRedis(unix_socket_path=UDS_PATH), "foobar")
with lock:
pytest.raises(RuntimeError, lock.acquire)
def lock(self, key, expire=None, id=None, auto_renewal=False):
return Lock(self.__client,
key,
expire=expire,
id=id,
auto_renewal=auto_renewal)
def test_signal_cleanup_on_reset(conn):
"""After resetting a lock, the signal key should not remain."""
lock = Lock(conn, 'foo')
lock.acquire()
lock.reset()
assert conn.llen('lock-signal:foo') == 0
def _execute_task_group(self, queue, tasks, all_task_ids):
"""
Executes the given tasks in the queue. Updates the heartbeat for task
IDs passed in all_task_ids. This internal method is only meant to be
called from within _process_from_queue.
"""
log = self.log.bind(queue=queue)
locks = []
# Keep track of the acquired locks: If two tasks in the list require
# the same lock we only acquire it once.
lock_ids = set()
ready_tasks = []
for task in tasks:
if task.get('lock', False):
if task.get('lock_key'):
kwargs = task.get('kwargs', {})
lock_id = gen_unique_id(
task['func'],
None,
{key: kwargs.get(key)
for key in task['lock_key']},
)
else:
lock_id = gen_unique_id(
task['func'],
task.get('args', []),
task.get('kwargs', {}),
)
if lock_id not in lock_ids:
lock = Lock(
self.connection,
self._key('lock', lock_id),
timeout=self.config['ACTIVE_TASK_UPDATE_TIMEOUT'])
acquired = lock.acquire(blocking=False)
if acquired:
lock_ids.add(lock_id)
locks.append(lock)
else:
log.info('could not acquire lock', task_id=task['id'])
# Reschedule the task
when = time.time() + self.config['LOCK_RETRY']
self._redis_move_task(queue, task['id'], ACTIVE,
SCHEDULED, when)
# Make sure to remove it from this list so we don't re-add
# to the ACTIVE queue by updating the heartbeat.
all_task_ids.remove(task['id'])
continue
ready_tasks.append(task)
if not ready_tasks:
return True, []
success = self._execute(queue, ready_tasks, log, locks, all_task_ids)
for lock in locks:
lock.release()
return success, ready_tasks
def handle(self, *args, **options):
with Lock(redis_connection, 'rescan_photos'):
self.rescan_photos(options['paths'])
def test_double_acquire(self):
lock = Lock(StrictRedis(unix_socket_path=UDS_PATH), "foobar")
with lock:
self.assertRaises(RuntimeError, lock.acquire)
def test_strict_check():
pytest.raises(ValueError, Lock, object(), name='foobar')
Lock(object(), name='foobar', strict=False)
def test_plain(redis_server):
with Lock(StrictRedis(unix_socket_path=UDS_PATH), "foobar"):
time.sleep(0.01)
def ensure_downloaded(self, lock_name=None):
if self.graph_cache_key in self.graph_cache:
return True
version_file = os.path.join(self.model_dir, self.name, 'version.txt')
if not lock_name:
lock_name = 'classifier_{}_download'.format(self.name)
r = redis.Redis(host=os.environ.get('REDIS_HOST', '127.0.0.1'))
with Lock(r, lock_name):
try:
with open(version_file) as f:
if f.read().strip() == str(self.version):
return True
except FileNotFoundError:
pass
model_info = self.get_model_info()
error = False
for file_data in model_info['files']:
final_path = os.path.join(self.model_dir, self.name,
file_data['filename'])
if not os.path.exists(final_path):
locations = file_data['locations']
index = random.choice(range(len(locations)))
location = locations.pop(index)
hash_sha256 = hashlib.sha256()
request = requests.get(location, stream=True)
if request.status_code != 200:
error = True
logger.error(f"Failed to fetch model for {location}: "
f"Got {request.status_code}")
continue
# Download file to temporary location
with tempfile.NamedTemporaryFile(mode='w+b',
delete=False) as f:
for chunk in request.iter_content(chunk_size=1024 *
1024): # 1MB chunks
if chunk: # filter out keep-alive new chunks
f.write(chunk)
hash_sha256.update(chunk)
# Move file to correct location if the hash matches
if hash_sha256.hexdigest() == file_data['sha256']:
dirname = os.path.dirname(final_path)
if not os.path.isdir(dirname):
os.makedirs(dirname)
if file_data.get('decompress'):
if file_data['filename'].endswith('.xz'):
xz_path = '{}.xz'.format(f.name)
shutil.move(f.name, xz_path)
subprocess.run(['unxz', xz_path])
final_path = final_path.replace('.xz', '')
shutil.move(f.name, final_path)
else:
error = True
logger.error(f"File downloaded from {location} is "
"corrupt as indicated by bad hash")
# TODO: Delete badly downloaded file
# Write version file
with open(version_file, 'w') as f:
if error:
f.write('ERROR\n')
return False
else:
f.write('{}\n'.format(str(self.version)))
return True
def test_reset(conn):
with Lock(conn, "foobar") as lock:
lock.reset()
new_lock = Lock(conn, "foobar")
new_lock.acquire(blocking=False)
new_lock.release()
def test_timeout_expire(conn):
lock1 = Lock(conn, "foobar", expire=1)
lock1.acquire()
lock2 = Lock(conn, "foobar")
assert lock2.acquire(timeout=2)
def test_bogus_release(conn):
lock = Lock(conn, "foobar-tok")
pytest.raises(NotAcquired, lock.release)
lock.release(force=True)
def test_not_usable_timeout(conn):
lock = Lock(conn, "foobar")
with pytest.raises(TimeoutNotUsable):
lock.acquire(blocking=False, timeout=1)
def test_release_from_nonblocking_leaving_garbage(conn):
for _ in range(10):
lock = Lock(conn, 'release_from_nonblocking')
lock.acquire(blocking=False)
lock.release()
assert conn.llen('lock-signal:release_from_nonblocking') == 1
def test_double_acquire(conn):
lock = Lock(conn, "foobar")
with lock:
pytest.raises(RuntimeError, lock.acquire)
pytest.raises(AlreadyAcquired, lock.acquire)
def test_auto_renewal_bad_values(conn):
with pytest.raises(ValueError):
Lock(conn, 'lock_renewal', expire=None, auto_renewal=True)
def test_owner_id(conn):
unique_identifier = "foobar-identifier"
lock = Lock(conn, "foobar-tok", expire=TIMEOUT / 4, id=unique_identifier)
lock_id = lock.id
assert lock_id == unique_identifier
def workerfn(unblocked):
conn = make_conn()
lock = Lock(conn, 'lock')
if lock.acquire():
unblocked.value = 1
def test_token(conn):
lock = Lock(conn, "foobar-tok")
tok = lock.id
assert conn.get(lock._name) is None
lock.acquire(blocking=False)
assert maybe_decode(conn.get(lock._name)) == tok
def workerfn(unblocked):
conn = make_conn()
lock1 = Lock(conn, 'lock1')
lock2 = Lock(conn, 'lock2')
if lock1.acquire() and lock2.acquire():
unblocked.value = 1
def test_no_auto_renewal(conn):
lock = Lock(conn, 'lock_renewal', expire=3, auto_renewal=False)
assert lock._lock_renewal_interval is None
lock.acquire()
assert lock._lock_renewal_thread is None, "No lock refresh thread should have been spawned"
if __name__ == '__main__':
if len(sys.argv) > 1 and sys.argv[1] == 'daemon':
logging.basicConfig(
level=logging.DEBUG,
format=
'%(process)d %(asctime)s,%(msecs)05d %(name)s %(levelname)s %(message)s',
datefmt="%x~%X")
test_name = sys.argv[2]
setup_coverage()
if test_name == 'test_simple':
conn = StrictRedis(unix_socket_path=UDS_PATH)
with Lock(conn, "foobar"):
time.sleep(0.1)
elif test_name == 'test_no_block':
conn = StrictRedis(unix_socket_path=UDS_PATH)
lock = Lock(conn, "foobar")
res = lock.acquire(blocking=False)
logging.info("acquire=>%s", res)
elif test_name == 'test_expire':
conn = StrictRedis(unix_socket_path=UDS_PATH)
with Lock(conn, "foobar", expire=TIMEOUT / 4):
time.sleep(0.1)
with Lock(conn, "foobar", expire=TIMEOUT / 4):
time.sleep(0.1)
elif test_name == 'test_no_overlap':
from sched import scheduler
sched = scheduler(time.time, time.sleep)
