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")