/java-api/src/net/stateupdateservice/StateService.java

package net.stateupdateservice;

import net.stateupdateservice.StateUpdateService.*;
import net.stateupdateservice.CacheEntry;
import net.stateupdateservice.ServiceMessage.ServicePacket;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.List;
import java.util.LinkedList;
import java.util.UUID;

import com.google.protobuf.*;


public class StateService {
	public int supportedPartialVersion = 2;

	public StateServiceServer service;
	public Map<String, WatchKey> watchKeys;
	private Map<String, CacheEntry> cache;

	private String currentWatchRequestSetId;
	private UpdateStateRequestSet.Builder currentUpdateStructure = null;
	private Set<String> sentDescriptorToServer = null;
	private String cacheFilePath = null;

	public StateService( StateServiceServer server, String cacheFile ) {
		this.service = server;
		this.cacheFilePath = cacheFile;
		watchKeys = new HashMap<String, WatchKey>();
		sentDescriptorToServer = new HashSet<String>();
		boolean loadedCache = false;

		if( cacheFile != null ) {
			File fhandle = new File( cacheFile );
			if( fhandle.exists() ) {
				try {
					FileInputStream fis = new FileInputStream( cacheFile );
					ObjectInputStream ois = new ObjectInputStream(fis);
					
					int cacheVersion = ois.readInt();
					if( cacheVersion == 1 ) {
						cache = (HashMap<String, CacheEntry>) ois.readObject();
						ByteString currentUpdateStructureString = (ByteString) ois.readObject();
						if( currentUpdateStructureString != null ) {
							currentUpdateStructure = UpdateStateRequestSet.newBuilder();
							currentUpdateStructure.mergeFrom(currentUpdateStructureString);
						}
						ois.close();
						loadedCache = true;
					}
				} catch ( IOException e) {
					
				} catch ( ClassNotFoundException e ) {
					
				}
			}
		}
		if(!loadedCache) {
			cache = new HashMap<String, CacheEntry>();
		}
	}

	public void flushCacheToDisk( ) {
		if( cacheFilePath != null ) {
			try {
				FileOutputStream fos = new FileOutputStream( cacheFilePath );
				ObjectOutputStream oos = new ObjectOutputStream(fos);
	
				oos.writeInt( 1 ); // Cache version 1
				oos.writeObject( cache );
				if( currentUpdateStructure != null ) {
					oos.writeObject( currentUpdateStructure.build().toByteString() );
				} else {
					oos.writeObject( null );
				}
				oos.close();
				
			} catch (IOException e) {
			
			}
		}
	}

	private class WatchKey {
		@SuppressWarnings("unused")
		public String key;
		public boolean allowPartial;

		public WatchKey( String key, boolean allowPartial ) {
			this.key = key;
			this.allowPartial = allowPartial;
		}
	}

	public class BadServiceResponseException extends IOException {
		/**
		 * 
		 */
		private static final long serialVersionUID = 1L;

		public BadServiceResponseException( String message ) {
			super( message );
		}
	}

	public StateService initKey( String key, Message.Builder valueObject ) {
		if( cache.containsKey(key) ) {
			// Make sure value is completely loaded
			cache.get(key).initValue(valueObject);
		} else {
			// Init specified keys to blank so we can operate on it
			Message.Builder obj = valueObject.clone();
			obj.clear();
	
			// Default to a generic not-valid entry.  This needs to be updated
			cache.put( key, new CacheEntry( key, obj, "", 0 ) );
		}
		return this;
	}

	public StateService updateKey( String key, Message.Builder msg ) {
		//  This queues an update into memory (and flushes to cache)
		if( currentUpdateStructure == null ) {
			currentUpdateStructure = UpdateStateRequestSet.newBuilder();
		}

		// Begin building update object
		UpdateStateRequestSet.UpdateStateRequest.Builder updateObj = 
				UpdateStateRequestSet.UpdateStateRequest.newBuilder()
							.setKey( key )
							.setValue( msg.build().toByteString() );

		// Add new update entry to structure
		if( !sentDescriptorToServer.contains( key ) ) {
			// Set the protobuf descriptor so the server can cache it
			updateObj.setProtobufDescriptor( msg.getDescriptorForType().toProto() );
		
			// Mark this so we don't need to send the descriptor every time
			sentDescriptorToServer.add( key );
		}
		currentUpdateStructure.addUpdates( updateObj );
		flushCacheToDisk();
		return this;
	}

	public StateService partialUpdateKey( String key, Message.Builder msg ) {
		if( currentUpdateStructure == null ) {
			currentUpdateStructure = UpdateStateRequestSet.newBuilder();
		}
		// Build update object
		currentUpdateStructure.addUpdates( UpdateStateRequestSet.UpdateStateRequest.newBuilder()
					.setKey( key )
					.setPartialUpdate( 
						PartialUpdate.newBuilder()
							.setType( PartialUpdate.PartialUpdateType.PROTOBUF_MERGE ) 
							.setAppendBytes( msg.buildPartial().toByteString() ) 
					)
		);
		return this;
	}

	public HashMap<String, Boolean> flushUpdates( ) throws BadServiceResponseException, IOException {
		//  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( currentUpdateStructure == null ) {
			return null;
		}
		// Build packet
		ServicePacket request = ServicePacket.newBuilder()
			.setPacketType( ServicePacket.ServicePacketType.UPDATE_STATE_REQUEST_SET)
			.setPacketData( currentUpdateStructure.build().toByteString() )
			.build();

		// Run request and get response
		ServicePacket response = service.sendRequest( request );

		// Make sure proper type was returned
		if( response.getPacketType() != ServicePacket.ServicePacketType.UPDATE_STATE_RESPONSE_SET ) {
			throw new BadServiceResponseException("Expected UPDATE_STATE_RESPONSE_SET packet");
		}

		// Unpack response
		UpdateStateResponseSet.Builder responseSet = UpdateStateResponseSet.newBuilder();
		try { 
			responseSet.mergeFrom( response.getPacketData() );
		} catch( InvalidProtocolBufferException e ) {
			throw new BadServiceResponseException("Received invalid UPDATE_STATE_RESPONSE_SET packet: " + e.toString() );
		}

		// Build a map, detailing what failed, if anything
		HashMap<String, Boolean> statusInfoMap = new HashMap<String, Boolean>(); 
		for( UpdateStateResponseSet.UpdateStateResponse responseObj: responseSet.getResponsesList() ) {
			if( !responseObj.getSuccess() ) {
				statusInfoMap.put( responseObj.getKey(), false );
			}
		}

		// If nothing failed, reset structure
		if( statusInfoMap.isEmpty() ) {
			currentUpdateStructure = null;
		}

		flushCacheToDisk();
		return statusInfoMap; // If this map has nothing in it, all updates were successful
	}

	private void updateCache( FetchStateResponseSet.FetchStateResponse responseObj ) throws BadServiceResponseException {
		
		if( responseObj.hasValue() ) {
			// Replace current value with new value
			try {
				cache.get( responseObj.getKey() ).setValueBytes( responseObj.getValue() );
			} catch( InvalidProtocolBufferException e ) {
				throw new BadServiceResponseException("Received invalid FETCH_STATE_RESPONSE_SET value field: " + e.toString() );
			}
		} else if( responseObj.hasPartialUpdate() ) {
			if( responseObj.getPartialUpdate().getType() == PartialUpdate.PartialUpdateType.PROTOBUF_MERGE ) {
				// Simple merge update
				try {
					cache.get( responseObj.getKey() ).mergeValueBytes( responseObj.getPartialUpdate().getAppendBytes() );
				} catch( InvalidProtocolBufferException e ) {
					throw new BadServiceResponseException("Received invalid FETCH_STATE_RESPONSE_SET partial update: " + e.toString() );
				}
			}
		}

		cache.get( responseObj.getKey() ).revision = responseObj.getRevision();
		cache.get( responseObj.getKey() ).session = responseObj.getSession();
	}

	public List<Message> getCachedKeys( String[] keys ) throws BadServiceResponseException, IOException {
		return getKeys( keys, true, true);
	}
	
	public List<Message> getKeys( String[] keys, boolean allowPartial ) throws BadServiceResponseException, IOException {
		return getKeys( keys, false, allowPartial);
	}
	
	private synchronized List<Message> getKeys( String[] keys, boolean cacheOnly, boolean allowPartial ) throws BadServiceResponseException, IOException {
		for( String key : keys ) {
			if( !cache.containsKey( key ) ) { // Make sure this key has been initialized, otherwise, raise error
				throw new IllegalArgumentException("Key must be initialized before being passed to getKeys");
			}
		}
		// Update keys unless we only want to go to cache
		if( !cacheOnly ) {
			// Return list of values corresponding to the specified keys
			FetchStateRequestSet.Builder rset = FetchStateRequestSet.newBuilder();

			for( String key : keys ) {
				FetchStateRequestSet.FetchStateRequest.Builder req = FetchStateRequestSet.FetchStateRequest.newBuilder()
						.setKey( key );
				if( cache.get( key ).revision > 0 ) {
					req.setSession( cache.get( key ).session );
					req.setRevision( cache.get( key ).revision );
					if( allowPartial ) {
						req.setPartialVersion( supportedPartialVersion );
					}
				}
				rset.addRequests( req.build() );
			}

			// No blocking for new updates, let's just get any new updates to keys
			rset.setBlocking( false );

			// Build packet
			ServicePacket request = ServicePacket.newBuilder()
				.setPacketType( ServicePacket.ServicePacketType.FETCH_STATE_REQUEST_SET)
				.setPacketData( rset.build().toByteString() )
				.build();

			// Run request and get response
			ServicePacket response;
			
			response = service.sendRequest( request );

			// Make sure proper type was returned
			if( response.getPacketType() != ServicePacket.ServicePacketType.FETCH_STATE_RESPONSE_SET ) {
				throw new BadServiceResponseException("Expected FETCH_STATE_RESPONSE_SET packet");
			}

			// Unpack response
			FetchStateResponseSet.Builder responseSet = FetchStateResponseSet.newBuilder();
			try { 
				responseSet.mergeFrom( response.getPacketData() );
			} catch( InvalidProtocolBufferException e ) {
				throw new BadServiceResponseException("Received invalid FETCH_STATE_RESPONSE_SET packet: " + e.toString() );
			}

			// Apply all updates to cache, if it's not here, either it didn't need any updates or was not found
			for( FetchStateResponseSet.FetchStateResponse responseObj: responseSet.getResponsesList() ) {
				if( cache.containsKey( responseObj.getKey() ) ) {  // Sanity check... server should NOT reply with something that we didn't ask for
					updateCache( responseObj );
				} else {
					throw new BadServiceResponseException("Server respoded with key that we didn't ask for");
				}
			}
		}

		// Now build responses
		List<Message> returnResponses = new LinkedList<Message>();
		for( String key : keys ) {
			if( cache.containsKey( key ) && cache.get( key ).revision > 0 ) { // Make sure this key has been initialized, otherwise, raise error
				returnResponses.add( cache.get( key ).getValue().build() );
			} else {
				// Key not found
				returnResponses.add( null );
			}
		}
		flushCacheToDisk();
		return returnResponses;
	}

	public StateService watchKey( String key, boolean allowPartial ) throws BadServiceResponseException, IOException {
		if( !cache.containsKey( key ) ) { // Make sure this key has been initialized, otherwise, raise error
			throw new IllegalArgumentException("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( !watchKeys.containsKey( key ) ) {
			watchKeys.put( key,  new WatchKey( key, allowPartial ) );
		}

		// Update the server with the news, if necessary
		if( currentWatchRequestSetId != null ) {
			FetchStateRequestSet.Builder rset = FetchStateRequestSet.newBuilder();

			// Add a single item to structure
			FetchStateRequestSet.FetchStateRequest.Builder req = FetchStateRequestSet.FetchStateRequest.newBuilder()
					.setKey( key );
			if( cache.get( key ).revision > 0 ) {
				req.setSession( cache.get( key ).session );
				req.setRevision( cache.get( key ).revision );
				if( allowPartial ) {
					req.setPartialVersion( supportedPartialVersion );
				}
			}
			rset.addRequests( req.build() );

			// Tell server we are ADDING this key to earlier request
			rset.setRequestSetModify( FetchStateRequestSet.RequestSetOperation.REQUEST_SET_ADD );
			// Don't block for new updates
			rset.setBlocking( false );
			rset.setRequestSetId( currentWatchRequestSetId );
			// Threadsafe TODO: Release add/remove lock

			// Update to server
			// Build packet
			ServicePacket request = ServicePacket.newBuilder()
				.setPacketType( ServicePacket.ServicePacketType.FETCH_STATE_REQUEST_SET)
				.setPacketData( rset.build().toByteString() )
				.build();

			// Run request and get response
			ServicePacket response = service.sendRequest( request );

			// Make sure proper type was returned
			if( response.getPacketType() != ServicePacket.ServicePacketType.FETCH_STATE_RESPONSE_SET ) {
				throw new BadServiceResponseException("Expected FETCH_STATE_RESPONSE_SET packet");
			}
		} else {
			// Threadsafe TODO: Release add/remove lock
		}
			
		return this;
	}

	public StateService unWatchKey( String key ) throws BadServiceResponseException, IOException {
		if( !cache.containsKey( key ) ) { // Make sure this key has been initialized, otherwise, raise error
			throw new IllegalArgumentException("Key must be initialized before being passed to unWatchKey");
		}

		// Remove key to watch list
		// Threadsafe TODO: Wait for add/remove lock, aquire it when possible
		if( watchKeys.containsKey( key ) ) {
			watchKeys.remove( key );
		}

		// Update the server with the news, if necessary
		if( currentWatchRequestSetId != null ) {
			FetchStateRequestSet.Builder rset = FetchStateRequestSet.newBuilder();

			// Add a single item to structure
			FetchStateRequestSet.FetchStateRequest.Builder req = FetchStateRequestSet.FetchStateRequest.newBuilder()
					.setKey( key );
			rset.addRequests( req.build() );

			// Tell server we are REMOVING this key to earlier request
			rset.setRequestSetModify( FetchStateRequestSet.RequestSetOperation.REQUEST_SET_REMOVE );
			// Don't block for new updates
			rset.setBlocking( false );
			rset.setRequestSetId( currentWatchRequestSetId );
			// Threadsafe TODO: Release add/remove lock

			// Update to server
			// Build packet
			ServicePacket request = ServicePacket.newBuilder()
				.setPacketType( ServicePacket.ServicePacketType.FETCH_STATE_REQUEST_SET)
				.setPacketData( rset.build().toByteString() )
				.build();

			// Run request and get response
			ServicePacket response = service.sendRequest( request );

			// Make sure proper type was returned
			if( response.getPacketType() != ServicePacket.ServicePacketType.FETCH_STATE_RESPONSE_SET ) {
				throw new BadServiceResponseException("Expected FETCH_STATE_RESPONSE_SET packet");
			}
		} else {
			// Threadsafe TODO: Release add/remove lock
		}
			
		return this;
	}

	public synchronized Map<String,Message> doWatch( ) throws BadServiceResponseException, IOException {
		// Block for updates on watched keys, then return map of updated keys that were being watched and changed

		// 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
		FetchStateRequestSet.Builder rset = FetchStateRequestSet.newBuilder();

		for( String key : watchKeys.keySet() ) {
			FetchStateRequestSet.FetchStateRequest.Builder req = FetchStateRequestSet.FetchStateRequest.newBuilder()
					.setKey( key );
			if( cache.get( key ).revision > 0 ) {
				req.setSession( cache.get( key ).session );
				req.setRevision( cache.get( key ).revision );
				if( watchKeys.get( key ).allowPartial ) {
					req.setPartialVersion( supportedPartialVersion );
				}
			}
			rset.addRequests( req.build() );
		}

		// Block for new updates
		rset.setBlocking( true );

		// Set ID for this requestset, so that we can amend it later without interrupting this send
		currentWatchRequestSetId = UUID.randomUUID().toString();
		rset.setRequestSetId( currentWatchRequestSetId );

		// Build packet
		ServicePacket request = ServicePacket.newBuilder()
			.setPacketType( ServicePacket.ServicePacketType.FETCH_STATE_REQUEST_SET)
			.setPacketData( rset.build().toByteString() )
			.build();

		// Threadsafe TODO: Release add/remove watch
		// This could block for a while, until new update comes
		// Run request and get response
		ServicePacket response = service.sendRequest( request );

		// 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
		currentWatchRequestSetId = null;
		// Threadsafe TODO: Release add/remove watch
		// Threadsafe TODO: Release watching

		// Make sure proper type was returned
		if( response.getPacketType() != ServicePacket.ServicePacketType.FETCH_STATE_RESPONSE_SET ) {
			throw new BadServiceResponseException("Expected FETCH_STATE_RESPONSE_SET packet");
		}

		// Unpack response
		FetchStateResponseSet.Builder responseSet = FetchStateResponseSet.newBuilder();
		try { 
			responseSet.mergeFrom( response.getPacketData() );
		} catch( InvalidProtocolBufferException e ) {
			throw new BadServiceResponseException("Received invalid FETCH_STATE_RESPONSE_SET packet: " + e.toString() );
		}

		// Apply all updates to cache
		// Threadsafe TODO: Lock add/remove watch, only for TRULY threaded languages
		Map<String, Message> returnResponses = new HashMap<String, Message>();

		for( FetchStateResponseSet.FetchStateResponse responseObj: responseSet.getResponsesList() ) {
			if( cache.containsKey( responseObj.getKey() ) ) {  // Sanity check... server should NOT reply with something that we didn't ask for
				updateCache( responseObj );

				// Add to the map of responses, if we have multiple updates for a single key, the old structure will be overwritten
				if( watchKeys.containsKey( responseObj.getKey() ) ) {
					returnResponses.put( responseObj.getKey(), cache.get( responseObj.getKey() ).getValue().build() );
				}
			} else {
				// Threadsafe TODO: Release add/remove watch, only for TRULY threaded languages
				throw new BadServiceResponseException("Server respoded with key that we didn't ask for");
			}
		}

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

		flushCacheToDisk();
		return returnResponses;
	}
}