def test_startswith(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in reversed(keys): pd[key] = None subset = [k for k in keys if k.startswith('ab')] self.assertSequenceEqual(subset, list(pd.startswith('ab')))
def test_slice_open_end(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in reversed(keys): pd[key] = key.upper() subset = [k.upper() for k in keys if not k.startswith('a')] self.assertSequenceEqual(subset, list(pd['b':]))
def test_slice_del_empty(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in keys: pd[key] = None del pd['e':] self.assertSequenceEqual(keys, list(pd))
def _getDirectoryEntriesFor( self, realm, category ): endpoints = {} with self.dirLock.reader(): cats = self.directory.get( realm, PrefixDict() )[ category : category ] for cat in cats: endpoints.update( cat ) return endpoints
def _updateDirectoryWith( self, curDir, newDir, newReverse ): ourNode = 'tcp://%s:' % ( self.ifaceIp4, ) isGhostActorsFound = False with self.dirLock.writer(): for uid, dest in newReverse.iteritems(): if uid not in self.reverseDir and uid not in self.tombstones: self.reverseDir[ uid ] = dest for realm, catMap in newDir.iteritems(): curDir.setdefault( realm, PrefixDict() ) for cat, endpoints in catMap.iteritems(): if 0 == len( endpoints ): continue curDir[ realm ].setdefault( cat, {} ) for uid, endpoint in endpoints.iteritems(): if uid in self.tombstones: continue # Check for ghost directory entries that report to be from here # but are not, may be that this node restarted. if endpoint.startswith( ourNode ) and uid not in self.actorInfo: self._addTombstone( uid ) isGhostActorsFound = True elif not endpoint.startswith( ourNode ): # Only add to this directory other node's info since # we are authoritative here. curDir[ realm ][ cat ][ uid ] = endpoint return curDir
def test_slice_set(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in reversed(keys): pd[key] = key newvalues = [k for k in keys if k.startswith('ab')] pd['ab':'ab'] = newvalues self.assertSequenceEqual(newvalues, list(pd['ab':'ab']))
def test_slice_set_short(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in reversed(keys): pd[key] = key newvalues = [k for k in keys if k.startswith('ab')] with self.assertRaises(ValueError): pd['ab':'ab'] = newvalues[:-1]
def test_slice_set_long(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in reversed(keys): pd[key] = key newvalues = [k for k in keys if k.startswith('ab')] pd['ab':'ab'] = newvalues + [None] self.assertNotIn(None, pd['ab':'ab'])
def test_slice_del(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in keys: pd[key] = None del pd['ab':'ab'] for key in keys: if not key.startswith('ab'): continue self.assertNotIn(key, pd)
def _svc_cleanupCats( self ): while not self.stopEvent.wait( 300 ): self._log( "Cleaning up directory" ) newDir = {} with self.dirLock.writer(): for realmName, realm in self.directory.iteritems(): newDir[ realmName ] = PrefixDict() for catName, cat in realm.iteritems(): if 0 != len( cat ): newDir[ realmName ][ catName ] = cat self.directory = newDir
def insert_search_delete(self, keys, value=None): pd = PrefixDict() for key in keys: val = key if value is None else value pd[key] = val self.assertEqual(len(pd), len(set(keys))) for key in keys: val = key if value is None else value self.assertIn(key, pd) self.assertEqual(pd[key], val) seen = set() for key in filterfalse(seen.__contains__, keys): seen.add(key) del pd[key] self.assertEqual(len(pd), 0) self.assertRaises(KeyError, operator.delitem, pd, '') for key in keys: self.assertRaises(KeyError, operator.getitem, pd, key) self.assertEqual(len(pd._root), 0)
def _svc_cleanupCats(self): while not self.stopEvent.wait(0): if self.isActorChanged.wait(2): # Buffer changes over 5 seconds gevent.sleep(5) self._log("Cleaning up directory") self.isActorChanged.clear() newDir = {} newNonOptDir = {} with self.dirLock.writer(): for realmName, realm in self.directory.iteritems(): newDir[realmName] = PrefixDict() newNonOptDir[realmName] = {} for catName, cat in realm.iteritems(): if 0 != len(cat): newDir[realmName][catName] = cat newNonOptDir[realmName][catName] = cat self.directory = newDir self.nonOptDir = newNonOptDir
def test_reversed(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in keys: pd[key] = None self.assertSequenceEqual(list(reversed(keys)), list(reversed(pd)))
def test_init_kwargs(self): pd = PrefixDict(a=None) self.assertIn('a', pd) self.assertIs(pd['a'], None)
def test_init_mapping(self): pd = PrefixDict({'a': None}) self.assertIn('a', pd) self.assertIs(pd['a'], None)
def test_slice_get_empty(self): pd = PrefixDict() keys = [''.join(combo) for combo in itertools.product('abc', repeat=3)] for key in reversed(keys): pd[key] = key.upper() self.assertSequenceEqual([], list(pd['d':'z']))
def test_sort_order(self): pd = PrefixDict() keys = ['', 'a', 'aa', 'ab', 'b', 'ba'] for key in reversed(keys): pd[key] = None self.assertSequenceEqual(keys, list(iter(pd)))
def test_commonprefix_full(self): pd = PrefixDict(abcd=None) self.assertEqual('abcd', pd.commonprefix('abcd'))
def test_commonprefix_empty(self): pd = PrefixDict(abcd=None) self.assertEqual(b'', pd.commonprefix('efgh'))
def commonsuffix(self, key): path = PrefixDict.commonprefix(self, key, restore_key=False) return path[::-1]
def prepare_key(self, key): return PrefixDict.prepare_key(self, key[::-1])
def restore_key(self, path, meta): key = PrefixDict.restore_key(self, path, meta) return key[::-1]
def test_commonprefix_half(self): pd = PrefixDict(abcd=None) self.assertEqual(b'ab', pd.commonprefix('abef'))
def test_slice_invalid_negative_step(self): pd = PrefixDict() with self.assertRaises(ValueError): pd[::-2]
def test_startswith_empty(self): pd = PrefixDict() pd['a'] = None self.assertSequenceEqual([], list(pd.startswith('b')))
def test_iter_post_del(self): pd = PrefixDict(a=0, b=1, c=2) del pd['b'] list(pd)
def test_invalid_key(self): pd = PrefixDict() self.assertRaises(TypeError, operator.setitem, pd, 0, None)
"""Test memory consumption and processing time with PrefixDict. Use words from '/usr/share/dict/words' as keys for PrefixDict and measure memory consumption and load time. """ import resource import sys from prefixtree import PrefixDict if __name__ == '__main__': start = resource.getrusage(resource.RUSAGE_SELF) glossary = PrefixDict() with open('/usr/share/dict/words') as words: for word in words: glossary[word.strip()] = None stop = resource.getrusage(resource.RUSAGE_SELF) rss_mb = (stop.ru_maxrss- start.ru_maxrss) / 1024.0 / 1024.0 tused = (stop.ru_utime + stop.ru_stime) sys.stdout.write('{0} MB\n'.format(rss_mb)) sys.stdout.write('{0} seconds\n'.format(tused))
def test_init_iterable(self): pd = PrefixDict([('a', None)]) self.assertIn('a', pd) self.assertIs(pd['a'], None)
def test_pickle(self): pd = PrefixDict() pd['a'] = None pickle.dumps(pd, pickle.HIGHEST_PROTOCOL)
def _svc_receiveOpsTasks( self ): z = self.opsSocket.getChild() while not self.stopEvent.wait( 0 ): data = z.recv() if data is not False and 'req' in data: action = data[ 'req' ] #self._log( "Received new ops request: %s" % action ) if 'keepalive' == action: z.send( successMessage() ) if 'from' in data and data[ 'from' ] not in self.nodes: self._log( "Discovered new node: %s" % data[ 'from' ] ) self._connectToNode( data[ 'from' ] ) for other in data.get( 'others', [] ): if other not in self.nodes: self._log( "Discovered new node: %s" % other ) self._connectToNode( other ) elif 'start_actor' == action: if not self._isPrivileged( data ): z.send( errorMessage( 'unprivileged' ) ) elif 'actor_name' not in data or 'cat' not in data: z.send( errorMessage( 'missing information to start actor' ) ) else: actorName = data[ 'actor_name' ] categories = data[ 'cat' ] realm = data.get( 'realm', 'global' ) parameters = data.get( 'parameters', {} ) resources = data.get( 'resources', {} ) ident = data.get( 'ident', None ) trusted = data.get( 'trusted', [] ) n_concurrent = data.get( 'n_concurrent', 1 ) owner = data.get( 'owner', None ) isIsolated = data.get( 'isolated', False ) log_level = data.get( 'loglevel', None ) log_dest = data.get( 'logdest', None ) uid = str( uuid.uuid4() ) port = self._getAvailablePortForUid( uid ) instance = self._getInstanceForActor( isIsolated ) if instance is not None: self._setActorMtd( uid, instance, actorName, realm, isIsolated, owner, parameters, resources ) newMsg = instance[ 'socket' ].request( { 'req' : 'start_actor', 'actor_name' : actorName, 'realm' : realm, 'uid' : uid, 'ip' : self.ifaceIp4, 'port' : port, 'parameters' : parameters, 'resources' : resources, 'ident' : ident, 'trusted' : trusted, 'n_concurrent' : n_concurrent, 'isolated' : isIsolated, 'loglevel' : log_level, 'logdest' : log_dest }, timeout = 30 ) else: newMsg = False if isMessageSuccess( newMsg ): self._log( "New actor loaded (isolation = %s, concurrent = %d), adding to directory" % ( isIsolated, n_concurrent ) ) # We always add a hardcoded special category _ACTORS/actorUid to provide a way for certain special actors # to talk to specific instances directly, but this is discouraged. with self.dirLock.writer(): self.reverseDir[ uid ] = 'tcp://%s:%d' % ( self.ifaceIp4, port ) self.directory.setdefault( realm, PrefixDict() ).setdefault( '_ACTORS/%s' % ( uid, ), {} )[ uid ] = 'tcp://%s:%d' % ( self.ifaceIp4, port ) for category in categories: self.directory.setdefault( realm, PrefixDict() ).setdefault( category, {} )[ uid ] = 'tcp://%s:%d' % ( self.ifaceIp4, port ) self.isActorChanged.set() else: self._logCritical( 'Error loading actor %s: %s.' % ( actorName, newMsg ) ) self._removeUidFromDirectory( uid ) self._addTombstone( uid ) z.send( newMsg ) elif 'kill_actor' == action: if not self._isPrivileged( data ): z.send( errorMessage( 'unprivileged' ) ) elif 'uid' not in data: z.send( errorMessage( 'missing information to stop actor' ) ) else: uids = data[ 'uid' ] if not isinstance( uids, collections.Iterable ): uids = ( uids, ) failed = [] for uid in uids: instance = self.actorInfo.get( uid, {} ).get( 'instance', None ) if instance is None: failed.append( errorMessage( 'actor not found' ) ) else: newMsg = instance[ 'socket' ].request( { 'req' : 'kill_actor', 'uid' : uid }, timeout = 20 ) if not isMessageSuccess( newMsg ): failed.append( newMsg ) if not self._removeUidFromDirectory( uid ): failed.append( errorMessage( 'error removing actor from directory after stop' ) ) self._addTombstone( uid ) self._removeInstanceIfIsolated( instance ) self.isActorChanged.set() if 0 != len( failed ): z.send( errorMessage( 'some actors failed to stop', failed ) ) else: z.send( successMessage() ) elif 'remove_actor' == action: if not self._isPrivileged( data ): z.send( errorMessage( 'unprivileged' ) ) elif 'uid' not in data: z.send( errorMessage( 'missing information to remove actor' ) ) else: uid = data[ 'uid' ] instance = self.actorInfo.get( uid, {} ).get( 'instance', None ) if instance is not None and self._removeUidFromDirectory( uid ): z.send( successMessage() ) self.isActorChanged.set() self._removeInstanceIfIsolated( instance ) else: z.send( errorMessage( 'actor to stop not found' ) ) elif 'host_info' == action: z.send( successMessage( { 'info' : { 'cpu' : psutil.cpu_percent( percpu = True, interval = 2 ), 'mem' : psutil.virtual_memory().percent } } ) ) elif 'get_full_dir' == action: with self.dirLock.reader(): z.send( successMessage( { 'realms' : self.directory, 'reverse' : self.reverseDir } ) ) elif 'get_dir' == action: realm = data.get( 'realm', 'global' ) if 'cat' in data: z.send( successMessage( data = { 'endpoints' : self._getDirectoryEntriesFor( realm, data[ 'cat' ] ) } ) ) else: z.send( errorMessage( 'no category specified' ) ) elif 'get_cats_under' == action: realm = data.get( 'realm', 'global' ) if 'cat' in data: with self.dirLock.reader(): z.send( successMessage( data = { 'categories' : [ x for x in self.directory.get( realm, PrefixDict() ).startswith( data[ 'cat' ] ) if x != data[ 'cat' ] ] } ) ) else: z.send( errorMessage( 'no category specified' ) ) elif 'get_nodes' == action: nodeList = {} for k in self.nodes.keys(): nodeList[ k ] = { 'last_seen' : self.nodes[ k ][ 'last_seen' ] } z.send( successMessage( { 'nodes' : nodeList } ) ) elif 'flush' == action: if not self._isPrivileged( data ): z.send( errorMessage( 'unprivileged' ) ) else: resp = successMessage() for uid, actor in self.actorInfo.items(): instance = actor[ 'instance' ] newMsg = instance[ 'socket' ].request( { 'req' : 'kill_actor', 'uid' : uid }, timeout = 30 ) if not isMessageSuccess( newMsg ): if newMsg is None or newMsg is False: self._log( 'stopping actor timed out, not all is lost' ) else: resp = errorMessage( 'error stopping actor' ) if not self._removeUidFromDirectory( uid ): if isMessageSuccess( resp ): resp = errorMessage( 'error removing actor from directory after stop' ) self._removeInstanceIfIsolated( instance ) z.send( resp ) if isMessageSuccess( resp ): self.isActorChanged.set() elif 'get_dir_sync' == action: with self.dirLock.reader(): z.send( successMessage( { 'directory' : self.directory, 'tombstones' : self.tombstones, 'reverse' : self.reverseDir } ) ) elif 'push_dir_sync' == action: if 'directory' in data and 'tombstones' in data and 'reverse' in data: z.send( successMessage() ) for uid, ts in data[ 'tombstones' ].iteritems(): self._addTombstone( uid, ts ) self._updateDirectoryWith( self.directory, data[ 'directory' ], data[ 'reverse' ] ) else: z.send( errorMessage( 'missing information to update directory' ) ) elif 'get_full_mtd' == action: z.send( successMessage( { 'mtd' : self.actorInfo } ) ) elif 'get_load_info' == action: info = {} for instance in self.processes: tmp = instance[ 'socket' ].request( { 'req' : 'get_load_info' }, timeout = 5 ) if isMessageSuccess( tmp ): info.update( tmp[ 'data' ] ) z.send( successMessage( { 'load' : info } ) ) elif 'associate' == action: if not self._isPrivileged( data ): z.send( errorMessage( 'unprivileged' ) ) else: uid = data[ 'uid' ] category = data[ 'category' ] try: info = self.actorInfo[ uid ] with self.dirLock.writer(): self.directory.setdefault( info[ 'realm' ], PrefixDict() ).setdefault( category, {} )[ uid ] = 'tcp://%s:%d' % ( self.ifaceIp4, info[ 'port' ] ) except: z.send( errorMessage( 'error associating, actor hosted here?' ) ) else: self.isActorChanged.set() z.send( successMessage() ) elif 'disassociate' == action: if not self._isPrivileged( data ): z.send( errorMessage( 'unprivileged' ) ) else: uid = data[ 'uid' ] category = data[ 'category' ] try: info = self.actorInfo[ uid ] with self.dirLock.writer(): del( self.directory[ info[ 'realm' ] ][ category ][ uid ] ) if 0 == len( self.directory[ info[ 'realm' ] ][ category ] ): del( self.directory[ info[ 'realm' ] ][ category ] ) except: z.send( errorMessage( 'error associating, actor exists in category?' ) ) else: self.isActorChanged.set() z.send( successMessage() ) else: z.send( errorMessage( 'unknown request', data = { 'req' : action } ) ) else: z.send( errorMessage( 'invalid request' ) ) self._logCritical( "Received completely invalid request" )