src/tashi/messaging/messaging.py - tashi - Git at Google

 #!/usr/bin/python

 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.

 import threading
 import thread
 import sys
 import os
 import socket
 import Queue
 import copy
 import random
 import traceback

 from threadpool import ThreadPoolClass, threadpool, ThreadPool
 from threadpool import threadpoolmethod, threaded, synchronized, synchronizedmethod

 class RWLock(object):
 	"""RWLock: Simple reader/writer lock implementation
 	FIXME: this implementation will starve writers!
 	Methods:
 		acquire() : take lock for read access
 		release() : release lock from read access
 		acquireWrite() : take lock for write access
 		releaseWrite() : release lock from write access"""
 	def __init__(self):
 		self.lock = threading.Condition()
 		self.readers = 0
 	def acquire(self):
 		self.lock.acquire()
 		self.readers = self.readers + 1
 		self.lock.release()
 	def release(self):
 		self.lock.acquire()
 		self.readers = self.readers - 1
 		self.lock.notify()
 		self.lock.release()
 	def acquireWrite(self):
 		self.lock.acquire()
 		while self.readers > 0:
 			self.lock.wait()
 	def releaseWrite(self):
 		self.lock.notify()
 		self.lock.release()


 class MessageBroker(object):
 	def __init__(self):
 		self.sublock = RWLock()
 		self.subscribers = []
 		self.random = random.Random()
 	def log(self, msg):
 		print "MessageBroker: Got log: '%s'" % str(msg)
 		return msg
 	def addSubscriber(self, subscriber):
 		self.sublock.acquireWrite()
 		self.subscribers.append(subscriber)
 		l = len(self.subscribers)
 		self.sublock.releaseWrite()
 		return l
 	def publish(self, message):
 		removesubs = []
 		i = self.random.randint(0,100)

 #		 subscribers = self.getSubscribers()
 #		 random.shuffle(subscribers)

 		self.sublock.acquire()

 		sys.stdout.flush()

 		for subscriber in self.subscribers:
 			try:
 				sys.stdout.flush()
 				assert(subscriber != self)
 				subscriber.publish(message)
 				sys.stdout.flush()
 			except Exception, e:
 				print e
 				removesubs.append(subscriber)

 		self.sublock.release()

 		if len(removesubs) > 0:
 			print "detected %i failed subscribers" % len(removesubs)
 			sys.stdout.flush()
 			self.sublock.acquireWrite()
 			for subscriber in removesubs:
 				try:
 					self.subscribers.remove(subscriber)
 				except:
 					pass
 			self.sublock.releaseWrite()
 	def getSubscribers(self):
 		self.sublock.acquire()
 		subs = copy.copy(self.subscribers)
 		self.sublock.release()
 		return subs
 	def removeSubscriber(self, subscriber):
 		self.sublock.acquireWrite()
 		try:
 			self.subscribers.remove(subscriber)
 		except:
 			pass
 		self.sublock.releaseWrite()
 	def publishList(self, messages):
 		for message in messages:
 			self.publish(message)

 class Subscriber(object):
 	def __init__(self, broker, pmatch={}, nmatch={}, synchronized=False):
 		self.broker = broker
 		self.lock = threading.Lock()
 		self.synchronized = synchronized
 		self.pmatch={}
 		self.nmatch={}
 		broker.addSubscriber(self)
 	def publish(self, message):
 		sys.stdout.flush()
 		msg = message
 		try:
 			if self.synchronized:
 				self.lock.acquire()
 			msg = self.filter(msg)
 			if (msg != None):
 				self.handle(msg)
 			if self.synchronized:
 				self.lock.release()
 		except Exception, x:
 			if self.synchronized:
 				self.lock.release()
 			print '%s, %s, %s' % (type(x), x, traceback.format_exc())
 	def publishList(self, messages):
 		for message in messages:
 			self.publish(message)
 	def handle(self, message):
 		print "Subscriber Default Handler: '%s'" % message
 	def setMatch(self, pmatch={}, nmatch={}):
 		self.lock.acquire()
 		self.pmatch=pmatch
 		self.nmatch=nmatch
 		self.lock.release()
 	def filter(self, message):
 		"""filter(self, message) : the filter function returns
 		the message, modified to be passed to the handler.
 		Returning (None) indicates that this is not a message
 		we are interested in, and it will not be passed to the
 		handler."""
 		send = True
 		for key in self.pmatch.keys():
 			if (not message.has_key(key)):
 				send = False
 				break
 			if self.pmatch[key] != None:
 				if message[key] != self.pmatch[key]:
 					send = False
 					break
 		if send == False:
 			return None
 		for key in message.keys():
 			if self.nmatch.has_key(key):
 				if self.nmatch[key] == None:
 					send = False
 					break
 				if self.nmatch[key] == message[key]:
 					send = False
 					break
 		if send == False:
 			return None
 		return message


 class Publisher(object):
 	'''Superclass for pub/sub publishers

 	FIXME: use finer-grained locking'''
 	def __init__(self, broker, aggregate=100):
 		self.pending = []
 		self.pendingLock = threading.Lock()
 		self.aggregateSize = aggregate
 		self.broker = broker
 	@synchronizedmethod
 	def publish(self, message):
 		if message.has_key('aggregate') and message['aggregate'] == 'True':
 			self.aggregate(message)
 			return
 		else:
 			self.broker.publish(message)
 	@synchronizedmethod
 	def publishList(self, messages):
 		self.broker.publishList(messages)
 	@synchronizedmethod
 	def aggregate(self, message):
 		# we can make this lock-less by using a queue for pending
 		# messages
 		self.pendingLock.acquire()
 		self.pending.append(message)
 		if len(self.pending) >= self.aggregateSize:
 			self.broker.publishList(self.pending)
 			self.pending = []
 		self.pendingLock.release()
 	@synchronizedmethod
 	def setBroker(self, broker):
 		self.broker = broker

 ##############################
 # Testing Code
 ##############################
 import time
 import unittest
 import sys
 import logging


 class TestSubscriber(Subscriber):
 	def __init__(self, *args, **kwargs):
 		self.queue = Queue.Queue()
 		Subscriber.__init__(self, *args, **kwargs)
 	def handle(self, message):
 		self.queue.put(message)

 class TestMessaging(unittest.TestCase):
 	def setUp(self):
 		self.broker = MessageBroker()
 		self.publisher = Publisher(self.broker)
 		self.subscriber = TestSubscriber(self.broker)
 	def testPublish(self):
 		self.publisher.publish( {'message':'hello world'} )
 		self.assertEqual(self.subscriber.queue.qsize(), 1)
 	def testPublishList(self):
 		nrmsgs = 10
 		msgs = []
 		for i in range(nrmsgs):
 			msgs.append( {'msgnum':str(i)} )
 		self.publisher.publishList( msgs )
 		self.assertEqual(self.subscriber.queue.qsize(), nrmsgs)
 	def testAggregate(self):
 		nrmsgs = self.publisher.aggregateSize
 		for i in range(nrmsgs):
 			self.assertEqual(self.subscriber.queue.qsize(), 0)
 			self.publisher.aggregate( {'msgnum':str(i)} )
 		self.assertEqual(self.subscriber.queue.qsize(), nrmsgs)
 	def testAggregateKeyword(self):
 		nrmsgs = self.publisher.aggregateSize
 		for i in range(nrmsgs):
 			self.assertEqual(self.subscriber.queue.qsize(), 0)
 			self.publisher.publish( {'msgnum':str(i), 'aggregate':'True'} )
 		self.assertEqual(self.subscriber.queue.qsize(), nrmsgs)

 if __name__ == '__main__':

 	logging.basicConfig(level=logging.INFO,
 						format="%(asctime)s %(levelname)s:\t %(message)s",
 						stream=sys.stdout)

 	suite = unittest.TestLoader().loadTestsFromTestCase(TestMessaging)
 	unittest.TextTestRunner(verbosity=2).run(suite)

 	sys.exit(0)


 ##############################
 # Old/Unused testing code
 ##############################


 	print 'testing removeSubscriber'
 	broker.removeSubscriber(subscriber)
 	publisher.publish( {'message':"you shouldn't see this"} )

 	nsub = NullSubscriber(broker)
 	print 'timing publish'
 	nrmsg = 100000
 	tt = time.time()
 	for i in range(nrmsg):
 #		publisher.publish( {"message":"hello world!"} )
 		publisher.publish( {} )
 	tt = time.time() - tt
 	print "Published %i messages in %f seconds, %f msg/s"%(nrmsg,
 														   tt,
 														   nrmsg/tt)
 	broker.removeSubscriber(nsub)

 	class SlowSubscriber(Subscriber):
 		def handle(self, message):
 			print 'called slow subscriber with message', message
 			time.sleep(1)
 			print 'returning from slow subscriber with message', message
 	class ThreadedSubscriber(Subscriber):
 		@threaded
 		def handle(self, message):
 			print 'called threaded subscriber with message', message
 			time.sleep(1)
 			print 'returning from threaded subscriber with message', message
 	class ThreadPoolSubscriber(Subscriber, ThreadPoolClass):
 		@threadpoolmethod
 		def handle(self, message):
 			print 'called threadpool subscriber with message', message
 			time.sleep(1)
 			print 'returning from threadpool subscriber with message', message


 	tsub = ThreadedSubscriber(broker)
 	for i in range(8):
 		publisher.publish( {"msg":str(i)} )
 	broker.removeSubscriber(tsub)
 	time.sleep(3)

 	tpsub = ThreadPoolSubscriber(broker)
 	for i in range(8):
 		publisher.publish( {"msg":str(i)} )
 	broker.removeSubscriber(tpsub)
 	time.sleep(3)

 	ssub = SlowSubscriber(broker)
 	for i in range(4):
 		publisher.publish( {"msg":str(i)} )
 	broker.removeSubscriber(ssub)
	#!/usr/bin/python

	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.

	import threading
	import thread
	import sys
	import os
	import socket
	import Queue
	import copy
	import random
	import traceback

	from threadpool import ThreadPoolClass, threadpool, ThreadPool
	from threadpool import threadpoolmethod, threaded, synchronized, synchronizedmethod

	class RWLock(object):
	"""RWLock: Simple reader/writer lock implementation
	FIXME: this implementation will starve writers!
	Methods:
	acquire() : take lock for read access
	release() : release lock from read access
	acquireWrite() : take lock for write access
	releaseWrite() : release lock from write access"""
	def __init__(self):
	self.lock = threading.Condition()
	self.readers = 0
	def acquire(self):
	self.lock.acquire()
	self.readers = self.readers + 1
	self.lock.release()
	def release(self):
	self.lock.acquire()
	self.readers = self.readers - 1
	self.lock.notify()
	self.lock.release()
	def acquireWrite(self):
	self.lock.acquire()
	while self.readers > 0:
	self.lock.wait()
	def releaseWrite(self):
	self.lock.notify()
	self.lock.release()



	class MessageBroker(object):
	def __init__(self):
	self.sublock = RWLock()
	self.subscribers = []
	self.random = random.Random()
	def log(self, msg):
	print "MessageBroker: Got log: '%s'" % str(msg)
	return msg
	def addSubscriber(self, subscriber):
	self.sublock.acquireWrite()
	self.subscribers.append(subscriber)
	l = len(self.subscribers)
	self.sublock.releaseWrite()
	return l
	def publish(self, message):
	removesubs = []
	i = self.random.randint(0,100)

	# subscribers = self.getSubscribers()
	# random.shuffle(subscribers)

	self.sublock.acquire()

	sys.stdout.flush()

	for subscriber in self.subscribers:
	try:
	sys.stdout.flush()
	assert(subscriber != self)
	subscriber.publish(message)
	sys.stdout.flush()
	except Exception, e:
	print e
	removesubs.append(subscriber)

	self.sublock.release()

	if len(removesubs) > 0:
	print "detected %i failed subscribers" % len(removesubs)
	sys.stdout.flush()
	self.sublock.acquireWrite()
	for subscriber in removesubs:
	try:
	self.subscribers.remove(subscriber)
	except:
	pass
	self.sublock.releaseWrite()
	def getSubscribers(self):
	self.sublock.acquire()
	subs = copy.copy(self.subscribers)
	self.sublock.release()
	return subs
	def removeSubscriber(self, subscriber):
	self.sublock.acquireWrite()
	try:
	self.subscribers.remove(subscriber)
	except:
	pass
	self.sublock.releaseWrite()
	def publishList(self, messages):
	for message in messages:
	self.publish(message)

	class Subscriber(object):
	def __init__(self, broker, pmatch={}, nmatch={}, synchronized=False):
	self.broker = broker
	self.lock = threading.Lock()
	self.synchronized = synchronized
	self.pmatch={}
	self.nmatch={}
	broker.addSubscriber(self)
	def publish(self, message):
	sys.stdout.flush()
	msg = message
	try:
	if self.synchronized:
	self.lock.acquire()
	msg = self.filter(msg)
	if (msg != None):
	self.handle(msg)
	if self.synchronized:
	self.lock.release()
	except Exception, x:
	if self.synchronized:
	self.lock.release()
	print '%s, %s, %s' % (type(x), x, traceback.format_exc())
	def publishList(self, messages):
	for message in messages:
	self.publish(message)
	def handle(self, message):
	print "Subscriber Default Handler: '%s'" % message
	def setMatch(self, pmatch={}, nmatch={}):
	self.lock.acquire()
	self.pmatch=pmatch
	self.nmatch=nmatch
	self.lock.release()
	def filter(self, message):
	"""filter(self, message) : the filter function returns
	the message, modified to be passed to the handler.
	Returning (None) indicates that this is not a message
	we are interested in, and it will not be passed to the
	handler."""
	send = True
	for key in self.pmatch.keys():
	if (not message.has_key(key)):
	send = False
	break
	if self.pmatch[key] != None:
	if message[key] != self.pmatch[key]:
	send = False
	break
	if send == False:
	return None
	for key in message.keys():
	if self.nmatch.has_key(key):
	if self.nmatch[key] == None:
	send = False
	break
	if self.nmatch[key] == message[key]:
	send = False
	break
	if send == False:
	return None
	return message



	class Publisher(object):
	'''Superclass for pub/sub publishers

	FIXME: use finer-grained locking'''
	def __init__(self, broker, aggregate=100):
	self.pending = []
	self.pendingLock = threading.Lock()
	self.aggregateSize = aggregate
	self.broker = broker
	@synchronizedmethod
	def publish(self, message):
	if message.has_key('aggregate') and message['aggregate'] == 'True':
	self.aggregate(message)
	return
	else:
	self.broker.publish(message)
	@synchronizedmethod
	def publishList(self, messages):
	self.broker.publishList(messages)
	@synchronizedmethod
	def aggregate(self, message):
	# we can make this lock-less by using a queue for pending
	# messages
	self.pendingLock.acquire()
	self.pending.append(message)
	if len(self.pending) >= self.aggregateSize:
	self.broker.publishList(self.pending)
	self.pending = []
	self.pendingLock.release()
	@synchronizedmethod
	def setBroker(self, broker):
	self.broker = broker

	##############################
	# Testing Code
	##############################
	import time
	import unittest
	import sys
	import logging


	class TestSubscriber(Subscriber):
	def __init__(self, args, *kwargs):
	self.queue = Queue.Queue()
	Subscriber.__init__(self, args, *kwargs)
	def handle(self, message):
	self.queue.put(message)

	class TestMessaging(unittest.TestCase):
	def setUp(self):
	self.broker = MessageBroker()
	self.publisher = Publisher(self.broker)
	self.subscriber = TestSubscriber(self.broker)
	def testPublish(self):
	self.publisher.publish( {'message':'hello world'} )
	self.assertEqual(self.subscriber.queue.qsize(), 1)
	def testPublishList(self):
	nrmsgs = 10
	msgs = []
	for i in range(nrmsgs):
	msgs.append( {'msgnum':str(i)} )
	self.publisher.publishList( msgs )
	self.assertEqual(self.subscriber.queue.qsize(), nrmsgs)
	def testAggregate(self):
	nrmsgs = self.publisher.aggregateSize
	for i in range(nrmsgs):
	self.assertEqual(self.subscriber.queue.qsize(), 0)
	self.publisher.aggregate( {'msgnum':str(i)} )
	self.assertEqual(self.subscriber.queue.qsize(), nrmsgs)
	def testAggregateKeyword(self):
	nrmsgs = self.publisher.aggregateSize
	for i in range(nrmsgs):
	self.assertEqual(self.subscriber.queue.qsize(), 0)
	self.publisher.publish( {'msgnum':str(i), 'aggregate':'True'} )
	self.assertEqual(self.subscriber.queue.qsize(), nrmsgs)

	if __name__ == '__main__':

	logging.basicConfig(level=logging.INFO,
	format="%(asctime)s %(levelname)s:\t %(message)s",
	stream=sys.stdout)

	suite = unittest.TestLoader().loadTestsFromTestCase(TestMessaging)
	unittest.TextTestRunner(verbosity=2).run(suite)

	sys.exit(0)


	##############################
	# Old/Unused testing code
	##############################



	print 'testing removeSubscriber'
	broker.removeSubscriber(subscriber)
	publisher.publish( {'message':"you shouldn't see this"} )

	nsub = NullSubscriber(broker)
	print 'timing publish'
	nrmsg = 100000
	tt = time.time()
	for i in range(nrmsg):
	# publisher.publish( {"message":"hello world!"} )
	publisher.publish( {} )
	tt = time.time() - tt
	print "Published %i messages in %f seconds, %f msg/s"%(nrmsg,
	tt,
	nrmsg/tt)
	broker.removeSubscriber(nsub)

	class SlowSubscriber(Subscriber):
	def handle(self, message):
	print 'called slow subscriber with message', message
	time.sleep(1)
	print 'returning from slow subscriber with message', message
	class ThreadedSubscriber(Subscriber):
	@threaded
	def handle(self, message):
	print 'called threaded subscriber with message', message
	time.sleep(1)
	print 'returning from threaded subscriber with message', message
	class ThreadPoolSubscriber(Subscriber, ThreadPoolClass):
	@threadpoolmethod
	def handle(self, message):
	print 'called threadpool subscriber with message', message
	time.sleep(1)
	print 'returning from threadpool subscriber with message', message



	tsub = ThreadedSubscriber(broker)
	for i in range(8):
	publisher.publish( {"msg":str(i)} )
	broker.removeSubscriber(tsub)
	time.sleep(3)

	tpsub = ThreadPoolSubscriber(broker)
	for i in range(8):
	publisher.publish( {"msg":str(i)} )
	broker.removeSubscriber(tpsub)
	time.sleep(3)

	ssub = SlowSubscriber(broker)
	for i in range(4):
	publisher.publish( {"msg":str(i)} )
	broker.removeSubscriber(ssub)