class Test(unittest.TestCase):
def setUp(self):
self.dir_ = tempfile.mkdtemp()
self.node = Store(self.dir_)
def tearDown(self):
self.node.shutdown()
shutil.rmtree(self.dir_)
def testLockExists(self):
self.assertFalse(os.path.exists(os.path.join(self.dir_, 'lock')))
def testPutGet(self):
get_num = lambda: random.randint(10000, 20000)
num1 = get_num()
self.assertEqual(self.node.put(str(num1)), str(num1))
num2 = get_num()
self.assertEqual(self.node.put(str(num2)), str(num2))
self.assertEqual(self.node.get(), str(num1))
self.assertEqual(self.node.get(), str(num2))
def testBatchPutGet(self):
self.node.shutdown()
size = 60
# '1'*20 will spend 30 bytes in `mq store`
batch1 = ['1' * 20, '2' * 20]
batch2 = ['3' * 20, '4' * 20]
self.node = Store(self.dir_, size)
self.assertEqual(self.node.put(batch1), batch1)
self.assertEqual(self.node.put(batch2), batch2)
self.assertEqual(len(self.node.legal_files), 2)
gets = sorted([self.node.get() for _ in range(4)])
res = list(batch1)
res.extend(batch2)
self.assertEqual(gets, res)
self.node.put('5' * 20)
self.assertEqual(self.node.get(), '5' * 20)
self.node.put('6' * 20)
self.assertEqual(len(self.node.legal_files), 1)
self.assertEqual(self.node.get(), '6' * 20)
self.assertEqual(self.node.get(), None)
self.assertRaises(StoreNoSpaceForPut, lambda: self.node.put('7' * 100))
def testPutCloseGet(self):
self.node.put('1' * 10)
self.node.shutdown()
self.node = Store(self.dir_)
self.assertEqual(self.node.get(), '1' * 10)
class LocalMessageQueueNode(object):
def __init__(self, base_dir, rpc_server, addr, addrs,
copies=1, n_priorities=3, deduper=None,
app_name=None):
self.dir_ = base_dir
self.rpc_server = rpc_server
assert addr in addrs
self.addr = addr
self.addrs = addrs
self.other_addrs = [n for n in self.addrs if n != self.addr]
self.copies = max(min(len(self.addrs)-1, copies), 0)
self.n_priorities = max(n_priorities, 1)
self.deduper = deduper
self.app_name = app_name
self._lock = threading.Lock()
self._register_rpc()
self.inited = False
def init(self):
with self._lock:
if self.inited: return
get_priority_store_dir = lambda priority: os.path.join(self.dir_,
PRIORITY_STORE_FN, str(priority))
self.priority_stores = [Store(get_priority_store_dir(i),
deduper=self.deduper,
mkdirs=True) \
for i in range(self.n_priorities)]
backup_store_dir = os.path.join(self.dir_, BACKUP_STORE_FN)
self.backup_stores = {}
for backup_addr in self.other_addrs:
backup_node_dir = backup_addr.replace(':', '_')
backup_path = os.path.join(backup_store_dir, backup_node_dir)
self.backup_stores[backup_addr] = Store(backup_path,
size=512*1024, mkdirs=True)
inc_store_dir = os.path.join(self.dir_, INCR_STORE_FN)
self.inc_store = Store(inc_store_dir, mkdirs=True)
self.inited = True
def _register_rpc(self):
if self.rpc_server:
self.register_rpc(self, self.rpc_server, app_name=self.app_name)
@classmethod
def register_rpc(cls, node, rpc_server, app_name=None):
prefix = get_rpc_prefix(app_name, 'mq')
rpc_server.register_function(node.put_proxy, name='put',
prefix=prefix)
rpc_server.register_function(node.batch_put_proxy, name='batch_put',
prefix=prefix)
rpc_server.register_function(node.put_backup_proxy, name='put_backup',
prefix=prefix)
rpc_server.register_function(node.get_proxy, name='get',
prefix=prefix)
rpc_server.register_function(node.exist, name='exist',
prefix=prefix)
def put(self, objs, force=False, priority=0):
self.init()
priority = max(min(priority, self.n_priorities-1), 0)
priority_store = self.priority_stores[priority]
priority_store.put(objs, force=force)
def put_proxy(self, pickled_objs, force=False, priority=0):
objs = pickle.loads(pickled_objs)
self.put(objs, force=force, priority=priority)
def batch_put(self, objs):
self.init()
puts = defaultdict(lambda:defaultdict(list))
for obj in objs:
priority = getattr(obj, 'priority', 0)
force = getattr(obj, 'force', False)
puts[priority][force].append(obj)
for priority, m in puts.iteritems():
for force, obs in m.iteritems():
self.put(obs, force=force, priority=priority)
def batch_put_proxy(self, pickled_objs):
objs = pickle.loads(pickled_objs)
self.batch_put(objs)
def put_backup(self, addr, objs, force=False):
self.init()
backup_store = self.backup_stores[addr]
backup_store.put(objs, force=force)
def put_backup_proxy(self, addr, pickled_objs, force=False):
objs = pickle.loads(pickled_objs)
self.put_backup(addr, objs, force=force)
def put_inc(self, objs, force=True):
self.init()
self.inc_store.put(objs, force=force)
def get(self, size=1, priority=0):
self.init()
priority = max(min(priority, self.n_priorities-1), 0)
priority_store = self.priority_stores[priority]
return priority_store.get(size=size)
def get_proxy(self, size=1, priority=0):
return pickle.dumps(self.get(size=size, priority=priority))
def get_backup(self, addr, size=1):
self.init()
backup_store = self.backup_stores[addr]
return backup_store.get(size=size)
def get_inc(self, size=1):
self.init()
return self.inc_store.get(size=size)
def add_node(self, addr):
if addr in self.addrs: return
self.addrs.append(addr)
backup_store_dir = os.path.join(self.dir_, BACKUP_STORE_FN)
backup_node_dir = addr.replace(':', '_')
backup_path = os.path.join(backup_store_dir, backup_node_dir)
self.backup_stores[addr] = Store(backup_path,
size=512*1024, mkdirs=True)
def remove_node(self, addr):
if addr not in self.addrs: return
self.addrs.remove(addr)
self.backup_stores[addr].shutdown()
del self.backup_stores[addr]
def exist(self, obj):
if self.deduper:
return self.deduper.exist(str(obj))
return False
def shutdown(self):
if not self.inited: return
[store.shutdown() for store in self.priority_stores]
for backup_store in self.backup_stores.values():
backup_store.shutdown()
self.inc_store.shutdown()
class LocalMessageQueueNode(object):
"""
Message queue node which only handle local mq operations
can also be in charge of handling the remote call.
This node includes several storage such as priority storage
for each priority, incremental storage as well as the
backup storage. Each storage is an instance of
:class:`~cola.core.mq.store.Store`.
"""
def __init__(self, base_dir, rpc_server, addr, addrs,
copies=1, n_priorities=3, deduper=None,
app_name=None):
self.dir_ = base_dir
self.rpc_server = rpc_server
assert addr in addrs
self.addr = addr
self.addrs = addrs
self.other_addrs = [n for n in self.addrs if n != self.addr]
self.copies = max(min(len(self.addrs)-1, copies), 0)
self.n_priorities = max(n_priorities, 1)
self.deduper = deduper
self.app_name = app_name
self._lock = threading.Lock()
self._register_rpc()
self.inited = False
def init(self):
with self._lock:
if self.inited: return
get_priority_store_dir = lambda priority: os.path.join(self.dir_,
PRIORITY_STORE_FN, str(priority))
self.priority_stores = [Store(get_priority_store_dir(i),
deduper=self.deduper,
mkdirs=True) \
for i in range(self.n_priorities)]
backup_store_dir = os.path.join(self.dir_, BACKUP_STORE_FN)
self.backup_stores = {}
for backup_addr in self.other_addrs:
backup_node_dir = backup_addr.replace(':', '_')
backup_path = os.path.join(backup_store_dir, backup_node_dir)
self.backup_stores[backup_addr] = Store(backup_path,
size=512*1024, mkdirs=True)
inc_store_dir = os.path.join(self.dir_, INCR_STORE_FN)
self.inc_store = Store(inc_store_dir, mkdirs=True)
self.inited = True
def _register_rpc(self):
if self.rpc_server:
self.register_rpc(self, self.rpc_server, app_name=self.app_name)
@classmethod
def register_rpc(cls, node, rpc_server, app_name=None):
prefix = get_rpc_prefix(app_name, 'mq')
rpc_server.register_function(node.put_proxy, name='put',
prefix=prefix)
rpc_server.register_function(node.batch_put_proxy, name='batch_put',
prefix=prefix)
rpc_server.register_function(node.put_backup_proxy, name='put_backup',
prefix=prefix)
rpc_server.register_function(node.get_proxy, name='get',
prefix=prefix)
rpc_server.register_function(node.exist, name='exist',
prefix=prefix)
def put(self, objs, force=False, priority=0):
self.init()
priority = max(min(priority, self.n_priorities-1), 0)
priority_store = self.priority_stores[priority]
priority_store.put(objs, force=force)
def put_proxy(self, pickled_objs, force=False, priority=0):
"""
The objects from remote call should be pickled to
avoid the serialization error.
:param pickled_objs: the pickled objects to put into mq
:param force: if set to True will be directly put into mq without
checking the duplication
:param priority: the priority queue to put into
"""
objs = pickle.loads(pickled_objs)
self.put(objs, force=force, priority=priority)
def batch_put(self, objs):
self.init()
puts = defaultdict(lambda:defaultdict(list))
for obj in objs:
priority = getattr(obj, 'priority', 0)
force = getattr(obj, 'force', False)
puts[priority][force].append(obj)
for priority, m in puts.iteritems():
for force, obs in m.iteritems():
self.put(obs, force=force, priority=priority)
def batch_put_proxy(self, pickled_objs):
"""
Unlike the :func:`put`, this method will check the ``priority``
of a single object to decide which priority queue to put into.
"""
objs = pickle.loads(pickled_objs)
self.batch_put(objs)
def put_backup(self, addr, objs, force=False):
self.init()
backup_store = self.backup_stores[addr]
backup_store.put(objs, force=force)
def put_backup_proxy(self, addr, pickled_objs, force=False):
"""
In the Cola backup mechanism, an object will not only be
put into a hash ring node, and also be put into the next
hash ring node which marked as a backup node. To the backup node,
it will remember the previous node's name.
:param addr: the node address to backup
:param pickled_objs: pickled objects
:param force: if True will be put into queue without checking duplication
"""
objs = pickle.loads(pickled_objs)
self.put_backup(addr, objs, force=force)
def put_inc(self, objs, force=True):
self.init()
self.inc_store.put(objs, force=force)
def get(self, size=1, priority=0):
self.init()
priority = max(min(priority, self.n_priorities-1), 0)
priority_store = self.priority_stores[priority]
return priority_store.get(size=size)
def get_proxy(self, size=1, priority=0):
"""
Get the objects from the specific priority queue.
:param size: if size == 1 will be the right object,
else will be the objects list
:param priority:
:return: unpickled objects
"""
return pickle.dumps(self.get(size=size, priority=priority))
def get_backup(self, addr, size=1):
self.init()
backup_store = self.backup_stores[addr]
return backup_store.get(size=size)
def get_inc(self, size=1):
self.init()
return self.inc_store.get(size=size)
def add_node(self, addr):
"""
When a new message queue node is in, firstly will add the address
to the known queue nodes, then a backup for this node will be created.
"""
if addr in self.addrs: return
self.addrs.append(addr)
backup_store_dir = os.path.join(self.dir_, BACKUP_STORE_FN)
backup_node_dir = addr.replace(':', '_')
backup_path = os.path.join(backup_store_dir, backup_node_dir)
self.backup_stores[addr] = Store(backup_path,
size=512*1024, mkdirs=True)
def remove_node(self, addr):
"""
For the removed node, this method is for the cleaning job including
shutting down the backup storage for the removed node.
"""
if addr not in self.addrs: return
self.addrs.remove(addr)
self.backup_stores[addr].shutdown()
del self.backup_stores[addr]
def exist(self, obj):
if self.deduper:
return self.deduper.exist(str(obj))
return False
def shutdown(self):
if not self.inited: return
[store.shutdown() for store in self.priority_stores]
for backup_store in self.backup_stores.values():
backup_store.shutdown()
self.inc_store.shutdown()
class Test(unittest.TestCase):
def setUp(self):
self.dir_ = tempfile.mkdtemp()
self.node = Store(self.dir_)
def tearDown(self):
self.node.shutdown()
shutil.rmtree(self.dir_)
def testLockExists(self):
self.assertFalse(os.path.exists(os.path.join(self.dir_, 'lock')))
def testPutGet(self):
get_num = lambda: random.randint(10000, 20000)
num1 = get_num()
self.assertEqual(self.node.put(str(num1)), str(num1))
num2 = get_num()
self.assertEqual(self.node.put(str(num2)), str(num2))
self.assertEqual(self.node.get(), str(num1))
self.assertEqual(self.node.get(), str(num2))
def testBatchPutGet(self):
self.node.shutdown()
size = 60
# '1'*20 will spend 30 bytes in `mq store`
batch1 = ['1' * 20, '2' * 20]
batch2 = ['3' * 20, '4' * 20]
self.node = Store(self.dir_, size)
self.assertEqual(self.node.put(batch1), batch1)
self.assertEqual(self.node.put(batch2), batch2)
self.assertEqual(len(self.node.legal_files), 2)
gets = sorted([self.node.get() for _ in range(4)])
res = list(batch1)
res.extend(batch2)
self.assertEqual(gets, res)
self.node.put('5' * 20)
self.assertEqual(self.node.get(), '5' * 20)
self.node.put('6' * 20)
self.assertEqual(len(self.node.legal_files), 1)
self.assertEqual(self.node.get(), '6' * 20)
self.assertEqual(self.node.get(), None)
self.assertRaises(StoreNoSpaceForPut, lambda: self.node.put('7' * 100))
def testPutCloseGet(self):
self.node.put('1'*10)
self.node.shutdown()
self.node = Store(self.dir_)
self.assertEqual(self.node.get(), '1'*10)