def test_zk_lost(self):
sess = {'acquired': True}
def watch(state):
dd('zk node state changed to: ', state)
sess['acquired'] = False
self.zk.add_listener(watch)
# test zk close
l = zkutil.ZKLock('foo_name', zkclient=self.zk)
l.acquire()
self.zk.stop()
time.sleep(0.1)
self.assertFalse(sess['acquired'])
# test node delete
sess['acquired'] = True
self.zk.start()
self.zk.add_listener(watch)
l = zkutil.ZKLock('foo_name', zkclient=self.zk)
with l:
time.sleep(0.1)
self.zk.delete(l.zkconf.lock('foo_name'))
time.sleep(0.1)
self.assertFalse(sess['acquired'])
def test_watch_acquire(self):
a = zkutil.ZKLock('foo', on_lost=lambda: True)
b = zkutil.ZKLock('foo', on_lost=lambda: True)
# no one locked
n = 0
for holder, ver in a.acquire_loop():
n += 1
self.assertEqual(0, n, 'acquired directly')
# watch node change
it = b.acquire_loop()
holder, ver = it.next()
self.assertEqual((a.identifier, 0), (holder, ver))
self.zk.set(a.lock_path, 'xx')
holder, ver = it.next()
self.assertEqual(('xx', 1), (holder, ver), 'watched node change')
a.release()
try:
holder, ver = it.next()
self.fail('should not have next yield')
except StopIteration:
pass
self.assertTrue(b.is_locked())
def test_node_change_after_acquired(self):
sess = {'acquired': True}
def on_lost():
sess['acquired'] = False
l = zkutil.ZKLock('foo_name',
zkclient=self.zk,
on_lost=on_lost)
with l:
sess['acquired'] = True
self.zk.delete(l.zkconf.lock('foo_name'))
time.sleep(0.1)
self.assertFalse(sess['acquired'])
l = zkutil.ZKLock('foo_name',
zkclient=self.zk,
on_lost=on_lost)
with l:
sess['acquired'] = True
self.zk.set(l.zkconf.lock('foo_name'), 'xxx')
time.sleep(0.1)
self.assertFalse(sess['acquired'])
def test_lock_holder(self):
a = zkutil.ZKLock('foo_name', on_lost=lambda: True)
b = zkutil.ZKLock('foo_name', on_lost=lambda: True)
with a:
self.assertEqual((a.identifier, 0), a.lock_holder)
val, zstate = self.zk.get(a.lock_path)
self.assertEqual((val, zstate.version), a.lock_holder)
locked, holder, ver = b.try_acquire()
self.assertFalse(locked)
self.assertEqual((a.identifier, 0), (holder, ver))
self.assertEqual((val, zstate.version), (holder, ver))
def test_try_release(self):
l1 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
l2 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
released, holder, ver = l1.try_release()
self.assertEqual((True, l1.identifier, -1), (released, holder, ver))
with l2:
released, holder, ver = l1.try_release()
self.assertEqual((False, l2.identifier, 0), (released, holder, ver))
released, holder, ver = l2.try_release()
self.assertEqual((True, l2.identifier, 0), (released, holder, ver))
def test_is_locked(self):
l = zkutil.ZKLock('foo_name', zkclient=self.zk)
with l:
pass
self.assertFalse(l.is_locked())
l = zkutil.ZKLock('foo_name', zkclient=self.zk)
l.acquire()
self.assertTrue(l.is_locked())
l.try_release()
self.assertFalse(l.is_locked())
def _loop_acquire(self, n, ident):
zk = KazooClient(hosts='127.0.0.1:21811')
zk.start()
scheme, name, passw = zk_test_auth
zk.add_auth(scheme, name + ':' + passw)
for ii in range(n):
l = zkutil.ZKLock('foo_name', zkclient=zk)
with l:
self.total += 1
self.counter += 1
self.assertTrue(self.counter == 1)
time.sleep(0.01)
self.counter -= 1
dd("id={ident:0>2} n={ii:0>2} got and released lock: {holder}".format(
ident=ident,
ii=ii,
holder=l.lock_holder))
zk.stop()
def begin(self, txid=None):
assert self.expire_at is None
self.expire_at = time.time() + self.timeout
if self.txid is None:
if txid is not None:
# Run a specified tx, normally when to recover a dead tx process
self.txid = txid
else:
# Run a new tx, create txid.
zstat = self.zke.set(self.zke._zkconf.txid_maker(), 'x')
self.txid = zstat.version
zkconf = copy.deepcopy(self.zke._zkconf.conf)
zkconf['lock_dir'] = self.zke._zkconf.tx_alive()
self.tx_alive_lock = zkutil.ZKLock(self.txid,
zkconf=zkconf,
zkclient=self.zke,
timeout=self.expire_at -
time.time())
try:
self.tx_alive_lock.acquire()
except zkutil.LockTimeout as e:
raise TXTimeout(repr(e) + ' while waiting for tx alive lock')
def test_config(self):
old = (config.zk_acl, config.zk_auth, config.zk_node_id)
config.zk_acl = (('foo', 'bar', 'cd'),
('xp', '123', 'cdrwa'))
config.zk_auth = ('digest', 'xp', '123')
config.zk_node_id = 'abc'
l = zkutil.ZKLock('foo_name', on_lost=lambda: True)
dd(l.zkconf.acl())
def _check_ac(ac):
self.assertEqual('digest', ac.id.scheme)
self.assertEqual('foo', ac.id.id.split(':')[0])
self.assertEqual(set(['CREATE', 'DELETE']), set(ac.acl_list))
_check_ac(l.zkconf.kazoo_digest_acl()[0])
with l:
# should have created lock node
data, zstate = self.zk.get(l.lock_path)
dd(data)
self.assertEqual('abc', data.split('-')[0])
acls, zstate = self.zk.get_acls(l.lock_path)
dd(acls)
_check_ac(acls[0])
(config.zk_acl, config.zk_auth, config.zk_node_id) = old
def setUp(self):
config.zk_acl = zk_test_acl
config.zk_auth = zk_test_auth
self.counter = 0
self.running = True
utdocker.create_network()
utdocker.start_container(
zk_test_name,
zk_test_tag,
env={
"ZOO_MY_ID": 1,
"ZOO_SERVERS": "server.1=0.0.0.0:2888:3888",
},
port_bindings={
2181: 21811,
}
)
self.zk = KazooClient(hosts='127.0.0.1:21811')
self.zk.start()
scheme, name, passw = zk_test_auth
self.zk.add_auth(scheme, name + ':' + passw)
# create lock base dir
acl = zkutil.make_kazoo_digest_acl(zk_test_acl)
self.zk.create('lock/', acl=acl)
self.lck = zkutil.ZKLock('foo_name', zkclient=self.zk)
def test_exception_and_recover(self):
# tx raised by other exception is recoverable
t = ZKTransaction(zkhost)
txid = None
try:
with ZKTransaction(zkhost) as t1:
start_ts = t1.start_ts
txid = t1.txid
f1 = t1.lock_get('foox')
f1.v = "foox_val"
t1.set(f1)
t1.set_state('bar')
start_ts = t1.start_ts
val, ver = t.zkstorage.state.get(txid)
self.assertEqual(
{
'got_keys': ["foox"],
'data': 'bar',
"start_ts": start_ts
}, val)
raise ValueError('foo')
except ValueError:
pass
val, ver = t.zkstorage.state.get(txid)
self.assertEqual(
{
'got_keys': ['foox'],
'data': 'bar',
"start_ts": start_ts
}, val)
ident = zkutil.make_identifier(txid, None)
keylock = zkutil.ZKLock('foox',
zkclient=t.zke,
zkconf=t.zke._zkconf,
ephemeral=False,
identifier=ident)
val, ver = keylock.get_lock_val()
self.assertEqual(val['v'], "foox_val")
with ZKTransaction(zkhost, txid=txid, timeout=1) as t1:
self.assertEqual(t1.start_ts, start_ts)
f1 = t1.lock_get('foox')
self.assertEqual(f1.v, "foox_val")
val = t1.get_state()
self.assertEqual(val, 'bar')
def _make_key_lock(self, txid, key):
keylock = zkutil.ZKLock(key,
zkclient=self.zke,
zkconf=self.zke._zkconf,
ephemeral=False,
identifier=utfjson.dump(txid))
return keylock
def test_lock_holder(self):
a = zkutil.ZKLock('foo_name', on_lost=lambda: True)
b = zkutil.ZKLock('foo_name', on_lost=lambda: True)
with a:
self.assertIsInstance(a.identifier, dict)
self.assertIsNone(a.identifier['val'], None)
self.assertEqual((a.identifier, 0), a.lock_holder)
val, zstate = self.zk.get(a.lock_path)
val = utfjson.load(val)
self.assertEqual((val, zstate.version), a.lock_holder)
locked, holder, ver = b.try_acquire()
self.assertFalse(locked)
self.assertEqual((a.identifier, 0), (holder, ver))
self.assertEqual((val, zstate.version), (holder, ver))
def test_persistent(self):
l = zkutil.ZKLock('foo_name', ephemeral=False, on_lost=lambda: True)
try:
with l:
l.zkclient.stop()
except ConnectionClosedError:
pass
self.assertRaises(zkutil.LockTimeout, self.lck.acquire, timeout=0.2)
def _make_key_lock(self, txid, key):
ident = zkutil.make_identifier(txid, None)
keylock = zkutil.ZKLock(key,
zkclient=self.zke,
zkconf=self.zke._zkconf,
ephemeral=False,
identifier=ident)
return keylock
def test_timeout(self):
l1 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
l2 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
with l1:
with ututil.Timer() as t:
self.assertRaises(zkutil.LockTimeout, l2.acquire, timeout=0.2)
self.assertAlmostEqual(0.2, t.spent(), places=1)
with ututil.Timer() as t:
self.assertRaises(zkutil.LockTimeout, l2.acquire, timeout=-1)
self.assertAlmostEqual(0.0, t.spent(), delta=0.01)
try:
l2.acquire(timeout=-1)
except zkutil.LockTimeout:
self.fail('timeout<0 should could acquire')
def test_release_owning_client_stopped(self):
l = zkutil.ZKLock('foo_name',
zkconf=dict(
hosts='127.0.0.1:21811',
),
on_lost=lambda: True)
l.release()
self.assertTrue(l.zkclient._stopped.is_set())
def test_specify_identifier(self):
a = zkutil.ZKLock('foo_name',
zkconf=dict(
hosts='127.0.0.1:21811',
),
identifier='faked',
on_lost=lambda: True)
b = zkutil.ZKLock('foo_name',
zkconf=dict(
hosts='127.0.0.1:21811',
),
identifier='faked',
on_lost=lambda: True)
a.acquire()
b.acquire()
dd('a and b has the same identifier thus they both can acquire the lock')
def test_hosts(self):
l = zkutil.ZKLock('foo_name',
zkconf=dict(
hosts='127.0.0.1:21811',
),
on_lost=lambda: True)
with l:
self.assertEqual('127.0.0.1:21811', l._hosts)
def begin(self, txid=None):
assert self.expire_at is None
self.expire_at = self.start_ts + self.timeout
if self.txid is None:
if txid is not None:
# Run a specified tx, normally when to recover a dead tx process
self.txid = txid
else:
# Run a new tx, create txid.
zstat = self.zke.set(self.zke._zkconf.txid_maker(), 'x')
self.txid = zstat.version
zkconf = copy.deepcopy(self.zke._zkconf.conf)
zkconf['lock_dir'] = self.zke._zkconf.tx_alive()
identifier = None
if self.zk_owner is not None:
identifier = {
'id': zkutil.lock_id(self.zke._zkconf.node_id()),
'val': {
'zk_owner': self.zk_owner
},
}
self.tx_alive_lock = zkutil.ZKLock(self.txid,
zkconf=zkconf,
zkclient=self.zke,
identifier=identifier,
timeout=self.expire_at -
time.time())
try:
self.tx_alive_lock.acquire()
except zkutil.LockTimeout as e:
raise TXTimeout(repr(e) + ' while waiting for tx alive lock')
state, ver = self._get_state(self.txid)
if state is not None:
if "start_ts" in state:
self.start_ts = state["start_ts"]
self.expire_at = self.start_ts + self.timeout
if 'got_keys' in state:
for key in state['got_keys']:
self.lock_get(key, latest=False)
if self.time_left() < 0:
raise TXTimeout('{tx} timeout when begin'.format(tx=self))
def test_try_lock(self):
l1 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
l2 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
with l1:
with ututil.Timer() as t:
locked, holder, ver = l2.try_acquire()
self.assertFalse(locked)
self.assertEqual(l1.identifier, holder)
self.assertGreaterEqual(ver, 0)
self.assertAlmostEqual(0.0, t.spent(), delta=0.05)
with ututil.Timer() as t:
locked, holder, ver = l2.try_acquire()
self.assertTrue(locked)
self.assertEqual(l2.identifier, holder)
self.assertEqual(ver, 0)
self.assertAlmostEqual(0.0, t.spent(), delta=0.05)
def test_conn_lost_when_blocking_acquiring(self):
l2 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
th = threadutil.start_daemon(target=self.zk.stop, after=0.5)
with l2:
try:
self.lck.acquire(timeout=1)
self.fail('expected connection error')
except ConnectionClosedError:
pass
th.join()
def test_internal_zkclient(self):
sess = {'acquired': True}
def on_lost():
sess['acquired'] = False
# There must be a listener specified to watch connection issue
self.assertRaises(ValueError, zkutil.ZKLock, 'foo_name')
l = zkutil.ZKLock('foo_name', on_lost=on_lost)
with l:
self.zk.delete(l.zkconf.lock('foo_name'))
time.sleep(0.1)
self.assertFalse(sess['acquired'])
def test_node_change_after_released(self):
sess = {'acquired': True}
def on_lost():
sess['acquired'] = False
l = zkutil.ZKLock('foo_name',
zkclient=self.zk,
on_lost=on_lost)
with l:
sess['acquired'] = True
time.sleep(0.1)
self.assertTrue(sess['acquired'])
def test_set_locked_key(self):
l1 = zkutil.ZKLock('foo_name', on_lost=lambda: True)
with l1:
val, zstate = self.zk.get(l1.lock_path)
val = utfjson.load(val)
self.assertEqual(val, l1.identifier)
self.assertEqual((val, zstate.version), l1.lock_holder)
self.assertIsInstance(val, dict)
self.assertIsNone(val['val'], None)
l1.set_lock_val("foo_val", zstate.version)
val, zstate = self.zk.get(l1.lock_path)
val = utfjson.load(val)
self.assertEqual(val, l1.identifier)
self.assertIsInstance(val, dict)
self.assertEqual(val['val'], "foo_val")
