/python-service/state_update_service.py

import zmq
from service_message_pb2 import *
from state_update_service_pb2 import *
import sys
import uuid

import traceback
from google.protobuf import reflection
from google.protobuf import descriptor
from google.protobuf import message

def MakeDescriptor(protobufDescriptor):
	newDescriptor = descriptor.MakeDescriptor( protobufDescriptor )
	for desc in protobufDescriptor.nested_type:
		nested_type = MakeDescriptor(desc) #descriptor.MakeDescriptor(desc)
		class concrete(message.Message):
			__metaclass__ = reflection.GeneratedProtocolMessageType
			DESCRIPTOR = nested_type
		nested_type._concrete_class = concrete					
		newDescriptor.nested_types.append(nested_type )
	newDescriptor.nested_types_by_name = dict((t.name, t) for t in newDescriptor.nested_types)
	
	for f in protobufDescriptor.field:
		type_name = f.type_name.split('.')[-1]
		if newDescriptor.nested_types_by_name.has_key(type_name):
			newDescriptor.fields_by_number[f.number].message_type = newDescriptor.nested_types_by_name[type_name]
	return newDescriptor

class StoreItem:
	def __init__( self, key, session ):
		self.key = key
		self.value = None
		# Current revision number of this store item
		self.revision = 0
		# The earliest revision number we can generate partial updates for
		self.revisionHistory = 0
		# The partial updates we have stored for this key
		self.history = []
		# Our current unique session key
		self.session = session
		# Listener addresses who are interested in this key
		self.listeners = {}
		# Set protobuf descriptor
		self.protobufDescriptor = None

	def log( self, message ):
		print "Log:",message

	def needsReply( self, requestObj ):
		if requestObj.HasField('session') and self.session != requestObj.session:
			return True
		if requestObj.HasField('revision') and self.revision == requestObj.revision:
			return False
		return True
		
	def updateValue( self, updateObj, replyItems ):
		# If somebody is trying to update this to a value that it is already at, just do nothing
		if not updateObj.HasField( 'partialUpdate' ) and self.value == updateObj.value:
			return True

		# Save any specified descriptor
		if updateObj.HasField( 'protobufDescriptor' ):
			try:
				self.protobufDescriptor = MakeDescriptor(updateObj.protobufDescriptor)
			except:
				traceback.print_exc(file=sys.stdout)
				return False

		# Do partial update, if it's sent
		if updateObj.HasField( 'partialUpdate' ):
			newValue = self.applyPartialUpdate( self.value, updateObj.partialUpdate )
			if newValue == None:
				return False

			self.value = newValue
			self.revision += 1
			# Leave revision history where it is, since we are able to build a partial update from there
			self.history.append( updateObj.partialUpdate )
		# This is a simple, update value
		elif updateObj.HasField( 'value' ):
			self.value = updateObj.value
			self.revision += 1
			self.revisionHistory = self.revision
			self.history = []

		else:
			# If they don't set a value OR a partial update, what are they expecting?
			self.log("Client didn't set value or partial update")
			return False
			
		# Add list of listeners to reply mapping
		for listener, originalRequest in self.listeners.items():
			if not replyItems.has_key( listener ):
				replyItems[ listener ] = []
			replyItems[ listener ].append( ( self, originalRequest ) )

		self.listeners = {}
		return True


	def applyPartialUpdate( self, value, partialUpdate ):
		if partialUpdate.type == PartialUpdate.APPEND and partialUpdate.HasField("appendBytes"):
			# Append the two messages
			return value + partialUpdate.appendBytes;

		# If we need to limit the number of fields, let's do it
		elif partialUpdate.type == PartialUpdate.PROTOBUF_MERGE and \
				partialUpdate.HasField("appendBytes") and self.protobufDescriptor != None:
			try:
				protoMsg = reflection.ParseMessage( self.protobufDescriptor, value )
				protoMsg.MergeFromString( partialUpdate.appendBytes )
			except:
				traceback.print_exc(file=sys.stdout)
				return None

			return protoMsg.SerializePartialToString()
		else:
			if not partialUpdate.HasField("appendBytes"):
				self.log("Client didn't set appendBytes field")
			elif self.protobufDescriptor == None:
				self.log("No protobuf descriptor has been set for this key: "+self.key)
			else:
				self.log("Client didn't specify any partial update operation")
			return None


	def addListener( self, listener, requestObj ):
		self.listeners[ listener ] = requestObj

	def removeListener( self, listener ):
		if self.listeners.has_key( listener ):
			del self.listeners[ listener ]

class StateUpdateService:
	def __init__( self, context, listenAddress ):
		self.context = context
		self.currentSession = str(uuid.uuid4())

		self.service = self.context.socket(zmq.ROUTER)
		self.service.bind( listenAddress )

		# TODO: Load from db
		self.kvstore = {}

	def log( self, message ):
		print "Log:",message

	def sendMessage( self, address, message ):
		self.service.send(address, zmq.SNDMORE)
		self.service.send('', zmq.SNDMORE)
		self.service.send( message )

	def flushListeners( self, replyItems ):
		# Go through each listener that we are replying to
		for listener in replyItems.keys():
			# Build response
			rep = FetchStateResponseSet()
			ind = 0
			for targetItem, requestObj in replyItems[ listener ]:

				# Assume, until proven otherwise that the history for the relevant key can support version 1 clients
				maxVersionNeededByHistory = 1
				if requestObj.HasField("partialVersion") and requestObj.partialVersion >= maxVersionNeededByHistory \
						and requestObj.HasField("revision") and requestObj.HasField("session") \
						and requestObj.session == targetItem.session \
						and requestObj.revision >= targetItem.revisionHistory \
						and requestObj.revision < targetItem.revision:
					# If we find a historical entry that needs version 2 of partial updates, we might need to bail
					for historicalEntry in targetItem.history[ requestObj.revision - targetItem.revisionHistory : ]:
						if historicalEntry.type not in [ PartialUpdate.APPEND, PartialUpdate.PROTOBUF_MERGE ]:
							maxVersionNeededByHistory = 2
				
				# Check whether this request is eligible for a partial update
				# 1. The partial version support is compatible
				# 2. Make sure we are talking about same session
				# 3. Make sure our history goes back far enough
				# 4. Make sure they are asking for a reasonable revision
				if requestObj.HasField("partialVersion") and requestObj.partialVersion >= maxVersionNeededByHistory \
						and requestObj.HasField("revision") and requestObj.HasField("session") \
						and requestObj.session == targetItem.session \
						and requestObj.revision >= targetItem.revisionHistory \
						and requestObj.revision < targetItem.revision:

					# Iterate over relevant history items and make responses for them
					for historicalEntry in targetItem.history[ requestObj.revision - targetItem.revisionHistory : ]:
						# Build fields for partial update
						rep.responses.add()
						# Set all required fields
						rep.responses[ ind ].session = targetItem.session
						rep.responses[ ind ].revision = targetItem.revision
						rep.responses[ ind ].key = targetItem.key
						# Copy partialUpdate object from history
						rep.responses[ ind ].partialUpdate.CopyFrom( historicalEntry )
					
						ind += 1
					
				# Otherwise, just send the value back, like normal
				else:
					rep.responses.add()
					# Set all required fields
					rep.responses[ ind ].session = targetItem.session
					rep.responses[ ind ].revision = targetItem.revision
					rep.responses[ ind ].key = targetItem.key
					# Set normal response
					rep.responses[ ind ].value = targetItem.value
					ind += 1

			# Build packet
			packet = ServicePacket()
			packet.packetType = ServicePacket.FETCH_STATE_RESPONSE_SET
			packet.packetData = rep.SerializeToString()

			# Send response packet
			self.sendMessage( listener, packet.SerializeToString() )
			
 	def main( self ):

		# Keep track of info for all request sets
		requestSetIdMap = {}
		reverseRequestSetIdMap = {}

		# Keep track of what entries have listeners
		listenerCacheMap = {}

		while True:
			[address, blank, message] = self.service.recv_multipart()

			try:
				packet = ServicePacket()
				packet.ParseFromString( message )
			except:
				traceback.print_exc(file=sys.stdout)
				continue

			# This is an update state message
			if packet.packetType == ServicePacket.UPDATE_STATE_REQUEST_SET:
				try:
					# Unserialize the request protobuf
					req = UpdateStateRequestSet()
					req.ParseFromString( packet.packetData )
				except:
					traceback.print_exc(file=sys.stdout)
					continue

				replyItems = {}
				# Build map of replies that need sent because of this update
				successes = []
				for requestObj in req.updates:
					# If this key doesn't exist yet, initialize new store item and update it
					if not self.kvstore.has_key( requestObj.key ):
						self.kvstore[ requestObj.key ] = StoreItem( requestObj.key, self.currentSession )
					# We must pass the reply map so it can add other listeners that need this update, if necessary
					successes.append( self.kvstore[ requestObj.key ].updateValue( requestObj, replyItems ) )

				# Build response
				rep = UpdateStateResponseSet()
				ind = 0
				for item in req.updates:
					rep.responses.add()
					rep.responses[ ind ].session = self.currentSession
					rep.responses[ ind ].revision = self.kvstore[ item.key ].revision
					rep.responses[ ind ].key = item.key
					rep.responses[ ind ].success = successes[ ind ]

					ind += 1

				# Build packet
				packet = ServicePacket()
				packet.packetType = ServicePacket.UPDATE_STATE_RESPONSE_SET
				packet.packetData = rep.SerializeToString()

				# Send response
				self.sendMessage( address, packet.SerializeToString() )

				# Now trigger all updates to be sent
				self.flushListeners( replyItems )

				for listener in replyItems.keys():
					# Cleanup hanging listeners
					if listenerCacheMap.has_key( listener ):
						for targetKey in listenerCacheMap[ listener ]:
							if self.kvstore.has_key( targetKey ):
								self.kvstore[ targetKey ].removeListener( listener )
						del listenerCacheMap[ listener ]

					# Cleanup requestSetIdMap
					if reverseRequestSetIdMap.has_key( listener ):
						if reverseRequestSetIdMap.has_key( listener ):
							del requestSetIdMap[ reverseRequestSetIdMap[ listener ] ]
							del reverseRequestSetIdMap[ listener ]

			# This is a fetch state message
			elif packet.packetType == ServicePacket.FETCH_STATE_REQUEST_SET:
				try:
					# Unserialize the request protobuf
					req = FetchStateRequestSet()
					req.ParseFromString( packet.packetData )
				except:
					traceback.print_exc(file=sys.stdout)
					continue

				if req.blocking:
					replyingNow = False
				else:
					replyingNow = True

				addingOperation = False
				removingOperation = False

				# If this is an operation on a previous requestset, let's overwrite the address, just in case we must reply
				if req.HasField("requestSetId"):
					if req.HasField("requestSetModify"):
						# Send blank response, since this is just a modify operation
						rep = FetchStateResponseSet()
						packet = ServicePacket()
						packet.packetType = ServicePacket.FETCH_STATE_RESPONSE_SET
						packet.packetData = rep.SerializeToString()

						# Send response packet
						self.sendMessage( address, packet.SerializeToString() )

						# If this is a well-formed modify packet, with a valid operation and a setId that hasn't been replied to yet, let's go forward
						if requestSetIdMap.has_key( req.requestSetId ) and req.requestSetModify in [ FetchStateRequestSet.REQUEST_SET_ADD, FetchStateRequestSet.REQUEST_SET_REMOVE ]:
							if req.requestSetModify == FetchStateRequestSet.REQUEST_SET_ADD:
								addingOperation = True
							else:
								removingOperation = True

							# Overwrite address with address from map
							address = requestSetIdMap[ req.requestSetId ]
							replyingNow = False # Force blocking
						else: # Something was badly formed, do nothing
							continue
					elif req.blocking:
						# Store this for future use, only if it's blocking request
						if requestSetIdMap.has_key( req.requestSetId ):
							# Somebody is sending the same requestsetid again, before first one has been replied to, we must cleanup old listener address
							oldListener = requestSetIdMap[ req.requestSetId ]
							if listenerCacheMap.has_key( oldListener ):
								for targetKey in listenerCacheMap[ oldListener ]:
									if self.kvstore.has_key( targetKey ):
										self.kvstore[ targetKey ].removeListener( oldListener )
								del listenerCacheMap[ oldListener ]

							# Send the old listener a blank reply, just to flush requests through the system
							blankrep = FetchStateResponseSet()
							blankpacket = ServicePacket()
							blankpacket.packetType = ServicePacket.FETCH_STATE_RESPONSE_SET
							blankpacket.packetData = blankrep.SerializeToString()
							self.sendMessage( oldListener, blankpacket.SerializeToString() )
							
						requestSetIdMap[ req.requestSetId ] = address
						reverseRequestSetIdMap[ address ] = req.requestSetId

				# Find out if we must reply now or later
				thisListenerRequests = { }
				replyItems = { address: [] }

				if not removingOperation:
					for item in req.requests:
							
						if self.kvstore.has_key( item.key ):
							thisListenerRequests[ item.key ] = item
							# Does this request need a reply right now?
							if self.kvstore[ item.key ].needsReply( item ):
								replyingNow = True  # In case it wasn't already set
								replyItems[ address ].append( ( self.kvstore[ item.key ], item ) )

				if replyingNow:
					# Send off all our replies
					self.flushListeners( replyItems )

					# On the off-chance that this is being triggered by a modify operation, let's cleanup any other keys being watched
					if listenerCacheMap.has_key( address ): 
						for targetKey in listenerCacheMap[ address ]:
							if self.kvstore.has_key( targetKey ):
								self.kvstore[ targetKey ].removeListener( address )
						del listenerCacheMap[ address ]
					# Cleanup requestSetIdMap
					if reverseRequestSetIdMap.has_key( address ):
						del requestSetIdMap[ reverseRequestSetIdMap[ address ] ]
						del reverseRequestSetIdMap[ address ]

				else:
					if removingOperation:
						# Ok, we are removing listeners, let's clean everything up
						for requestObj in req.requests:
							if self.kvstore.has_key( requestObj.key ):
								self.kvstore[ requestObj.key ].removeListener( address )

								if listenerCacheMap.has_key( address ):
									if requestObj.key in listenerCacheMap[ address ]:
										listenerCacheMap[ address ].remove( requestObj.key ) 
					else:
						# If we're not replying now, store this request as a listener on each of the keys he wants updates on
						if not listenerCacheMap.has_key( address ):
							listenerCacheMap[ address ] = []
						for requestObj in req.requests:
							if self.kvstore.has_key( requestObj.key ):
								self.kvstore[ requestObj.key ].addListener( address, requestObj )
								# Save in cache for later cleanup
								listenerCacheMap[ address ].append( requestObj.key ) 


if __name__ == "__main__":
	context   = zmq.Context(1)
	service = StateUpdateService( context, sys.argv[1] )
	service.main( )