/python-api/state_update_service_api.py

import sys
import copy
import uuid
from google.protobuf.descriptor import FieldDescriptor


sys.path.append('../python-service')

from service_message_pb2 import *
from state_update_service_pb2 import *
from descriptor_pb2 import DescriptorProto
import urllib2

def sendRequest( service, packetType, packetData ):
	# Build packet structure
	p = ServicePacket()
	p.packetType = packetType
	p.packetData = packetData
	# Send to backend
	service.send( p.SerializeToString() )
	# Get response, parse back into ServicePAcket
	p.ParseFromString( service.recv() )
	# Return type, data
	return p.packetType, p.packetData

class CacheEntry:
	def __init__( self, key, value, session, revision ):
		self.key = key
		self.value = value
		self.session = session
		self.revision = revision

class HttpService:
	def __init__( self, url ):
		self.url = url
		self.response = None

	def send( self, data ):
		req = urllib2.Request( self.url, data )
		req.add_header('Content-Type', 'application/octet-stream')
		while True:
			try:
				handle = urllib2.urlopen( req, timeout = 300 )
			except IOError:
				continue
			break
		self.response = handle.read()
		handle.close()

	def recv( self ):
		if self.response:
			return self.response
		else:
			raise IOError("You must successfully send before trying to receive")

class StateService:
	supportedPartialVersion = 1  # This client can deal with partial update version 2 feature set

	def __init__( self, context, server ):
		if server.startswith('tcp') or server.startswith('ipc'):
			import zmq
			self.service = context.socket( zmq.REQ )
			self.service.connect( server )

			self.watcher = context.socket( zmq.REQ )
			self.watcher.connect( server )

		elif server.startswith('http'):
			self.service = HttpService( server )
			self.watcher = self.service
		
		self.watchKeys = {}
		self.lastSession = None
		self.cache = {}
		self.currentWatchRequestSetId = None
		self.currentUpdateStructure = None
		self.sentDescriptorToServer = {}

	def initKey( self, key, valueObject ):
		# Init specified key to blank so we can operate on it
		obj = copy.deepcopy( valueObject )
		obj.Clear()
		# Default to a generic not-valid entry.  This needs to be updated
		self.cache[ key ] = CacheEntry( key, obj, "" , 0 )
		return self

	def updateKey( self, key, msg ):
		#  This queues an update into memory (and flushes to cache)
		if not self.currentUpdateStructure:
			self.currentUpdateStructure = UpdateStateRequestSet()

		# Add new update entry to structure
		self.currentUpdateStructure.updates.add()
		self.currentUpdateStructure.updates[ -1 ].key = key
		self.currentUpdateStructure.updates[ -1 ].value = msg.SerializeToString()

		if not self.sentDescriptorToServer.has_key( key ):
			msg.DESCRIPTOR.CopyToProto( self.currentUpdateStructure.updates[ -1 ].protobufDescriptor )
			# Mark this so we don't need to send the descriptor every time
			self.sentDescriptorToServer[ key ] = True

	def partialUpdateKey( self, key, msg ):

		if not self.currentUpdateStructure:
			self.currentUpdateStructure = UpdateStateRequestSet()

		# Add new partial update entry to structure
		self.currentUpdateStructure.updates.add()
		self.currentUpdateStructure.updates[ -1 ].key = key
		self.currentUpdateStructure.updates[ -1 ].partialUpdate.type = PartialUpdate.PROTOBUF_MERGE
		self.currentUpdateStructure.updates[ -1 ].partialUpdate.appendBytes = msg.SerializePartialToString()

	def flushUpdates( self ):
		#  Flush any state updates in queue to server (maybe loaded from cache, if we've crashed)
		# returns map of what updates failed, empty map otherwise

		# Check whether there is anything to flush
		if self.currentUpdateStructure == None:
			return {}
		repType, rep = sendRequest( self.service, ServicePacket.UPDATE_STATE_REQUEST_SET, self.currentUpdateStructure.SerializeToString() )
		assert repType == ServicePacket.UPDATE_STATE_RESPONSE_SET
		res = UpdateStateResponseSet()
		res.ParseFromString( rep )

		statusInfoMap = {} # Build a map, detailing what failed, if anything
		for response in res.responses:
			if not response.success:
				statusInfoMap[ response.key ] = False

		# If nothing failed, reset structure
		if len(statusInfoMap.keys()) == 0:
			self.currentUpdateStructure = None

		return statusInfoMap # If this map has nothing in it, all updates were successful

	def _updateCache( self, responseObj ):
		
		if responseObj.value:
			# Replace current value with new value
			self.cache[ responseObj.key ].value.ParseFromString( responseObj.value )

		elif responseObj.partialUpdate:
			if responseObj.partialUpdate.type == PartialUpdate.PROTOBUF_MERGE:
				# Simple merge update
				self.cache[ responseObj.key ].value.MergeFromString( responseObj.partialUpdate.appendBytes )

		self.cache[ responseObj.key ].revision = responseObj.revision
		self.cache[ responseObj.key ].session = responseObj.session
				

	def getKeys( self, keys, cacheOnly = False, allowPartial = True ):
		for ind, key in enumerate(keys):
			if not self.cache.has_key( key ): # Make sure this key has been initialized, otherwise, raise error
				raise Exception("Key must be initialized before being passed to getKeys")

		# Update keys unless we only want to go to cache
		if not cacheOnly:
			# Return list of values corresponding to the specified keys
			rset = FetchStateRequestSet()

			for ind, key in enumerate(keys):
				rset.requests.add()
				rset.requests[ ind ].key = key
				if self.cache[ key ].revision > 0:
					rset.requests[ ind ].session = self.cache[ key ].session
					rset.requests[ ind ].revision = self.cache[ key ].revision
					if allowPartial:
						rset.requests[ ind ].partialVersion = self.supportedPartialVersion

			# No blocking for new updates, let's just get any new updates to keys
			rset.blocking = False

			repType, rep = sendRequest( self.service, ServicePacket.FETCH_STATE_REQUEST_SET, rset.SerializeToString() )
			assert repType == ServicePacket.FETCH_STATE_RESPONSE_SET
			res = FetchStateResponseSet()
			res.ParseFromString( rep )

			# Apply all updates to cache, if it's not here, either it didn't need any updates or was not found
			for response in res.responses:
				if self.cache.has_key( response.key ): # Sanity check... server should NOT reply with something that we didn't ask for
					self._updateCache( response )
				else:
					raise Exception("Server respoded with key that we didn't ask for")

		# Now build responses
		responses = []
		for ind, key in enumerate(keys):
			if self.cache.has_key( key ) and self.cache[ key ].revision > 0:
				responses.append( copy.deepcopy( self.cache[ key ].value ) )
			else:
				# Key not found
				responses.append( None )
		return responses

	def watchKey( self, key, allowPartial = True ):
		if not self.cache.has_key( key ): # Make sure this key has been initialized, otherwise, raise error
			raise Exception("Key must be initialized before being passed to watchKey")

		# Add key to watch list
		# Threadsafe TODO: Wait for add/remove lock, aquire it when possible
		if not key in self.watchKeys.keys():
			self.watchKeys[ key ] = { "key": key, "allowPartial" : allowPartial }

		# Update the server with the news, if necessary
		if self.currentWatchRequestSetId != None:
			rset = FetchStateRequestSet()
			rset.requests.add()
			rset.requests[ 0 ].key = key
			if self.cache[ key ].revision > 0:
				rset.requests[ 0 ].session = self.cache[ key ].session
				rset.requests[ 0 ].revision = self.cache[ key ].revision
				if allowPartial:
					rset.requests[ 0 ].partialVersion = self.supportedPartialVersion

			# Tell server we are ADDING this key to earlier request
			rset.requestSetModify = FetchStateRequestSet.REQUEST_SET_ADD
			# Don't block for new updates
			rset.blocking = False
			rset.requestSetId = self.currentWatchRequestSetId
			# Threadsafe TODO: Release add/remove lock

			# Update to server
			repType, rep = sendRequest( self.service, ServicePacket.FETCH_STATE_REQUEST_SET, rset.SerializeToString() )
			assert repType == ServicePacket.FETCH_STATE_RESPONSE_SET
		else:
			pass
			# Threadsafe TODO: Release add/remove lock
			
		return self

	def unWatchKey( self, key ):
		if not self.cache.has_key( key ): # Make sure this key has been initialized, otherwise, raise error
			raise Exception("Key must be initialized before being passed to unWatchKey")

		# Threadsafe TODO: Wait for add/remove lock, aquire it when possible
		# Remove key from watch list
		if key in self.watchKeys.keys():
			del self.watchKeys[ key ]

		# Update the server with the news, if necessary
		if self.currentWatchRequestSetId != None:
			rset = FetchStateRequestSet()
			rset.requests.add()
			rset.requests[ 0 ].key = key
			# Tell server we are REMOVING this key from earlier request
			rset.requestSetModify = FetchStateRequestSet.REQUEST_SET_REMOVE
			# Don't block for new updates
			rset.blocking = False
			rset.requestSetId = self.currentWatchRequestSetId
			# Threadsafe TODO: Release add/remove lock

			# Update to server
			repType, rep = sendRequest( self.service, ServicePacket.FETCH_STATE_REQUEST_SET, rset.SerializeToString() )
			assert repType == ServicePacket.FETCH_STATE_RESPONSE_SET
		else:
			pass
			# Threadsafe TODO: Release add/remove lock

		return self

	def doWatch( self ):
		# Block for updates on watched keys, then return map of updated keys that were being watched and changed
		# Return list of values corresponding to the specified keys

		# Threadsafe TODO: Lock watching, fail if we are already watching
		# Threadsafe TODO: Lock add/remove watch, so we don't allow other threads to add/remove before it is safe
		rset = FetchStateRequestSet()

		for ind, key in enumerate( self.watchKeys.keys() ):
			rset.requests.add()
			rset.requests[ ind ].key = key
			if self.cache[ key ].revision > 0:
				rset.requests[ ind ].session = self.cache[ key ].session
				rset.requests[ ind ].revision = self.cache[ key ].revision
				if self.watchKeys[ key ]["allowPartial"]:
					rset.requests[ ind ].partialVersion = self.supportedPartialVersion

		# Block for new updates
		rset.blocking = True

		# Set ID for this requestset, so that we can amend it later without interrupting this send
		rset.requestSetId = str( uuid.uuid4() )
		self.currentWatchRequestSetId = rset.requestSetId

		# Threadsafe TODO: Release add/remove watch
		# This could block for a while, until new update comes
		repType, rep = sendRequest( self.watcher, ServicePacket.FETCH_STATE_REQUEST_SET, rset.SerializeToString() )

		# Threadsafe TODO: Lock add/remove watch, so we don't change current watch request in the middle of an add/remove operation
		# Now unset this, so other threads don't think we still have this request going
		self.currentWatchRequestSetId = None
		# Threadsafe TODO: Release add/remove watch
		# Threadsafe TODO: Release watching
		assert repType == ServicePacket.FETCH_STATE_RESPONSE_SET

		res = FetchStateResponseSet()
		res.ParseFromString( rep )

		# Apply all updates to cache
		responses = {}

		# Threadsafe TODO: Lock add/remove watch, only for TRULY threaded languages
		for response in res.responses:
			if self.cache.has_key( response.key ): # Sanity check... server should NOT reply with something that we didn't ask for
				self._updateCache( response )
				# Add to the map of responses, if we have multiple updates for a single key, the old structure will be overwritten
				if self.watchKeys.has_key( response.key ):
					responses[ response.key ] = copy.deepcopy( self.cache[ response.key ].value )
			else:
				# Threadsafe TODO: Release add/remove watch, only for TRULY threaded languages
				raise Exception("Server respoded with key that we didn't ask for")

		# Threadsafe TODO: Release add/remove watch, only for TRULY threaded languages

		return responses