content: /**
   * 
   */
  package com.sun.faban.driver.transport.asynchronous;
  
  import java.io.InputStream;
  import java.lang.management.ManagementFactory;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.Set;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.logging.Level;
  
  import javax.jms.Message;
  import javax.management.MBeanServer;
  
  import net.sf.ehcache.Cache;
  import net.sf.ehcache.CacheManager;
  import net.sf.ehcache.Element;
  import net.sf.ehcache.management.ManagementService;
  
  import org.apache.commons.math.stat.Frequency;
  import org.apache.commons.math.stat.descriptive.DescriptiveStatistics;
  
  import com.sun.faban.driver.DriverContext;
  
  /**
   * @author Noah Campbell
   * @param <T> result
   * @param <K> message
   *
   */
  public class MessageCacheTraceRegistry<T, K extends Message> implements TraceRegistry<T, K> {
  	
  	static private AtomicLong traceCounter = new AtomicLong(100000);
  	static private ScheduledExecutorService scheduler;
  	static private DescriptiveStatistics cacheStatistics = DescriptiveStatistics.newInstance();
  	static private Frequency roundtripFrequence = new Frequency();	
  	static private CacheManager responseCacheManager;
  	
  	private Cache cache;
  	private DriverContext context;
  	
  	/**
  	 * @param ctx
  	 * @throws RegistryConfigurationException 
  	 */
  	public MessageCacheTraceRegistry(DriverContext ctx) throws RegistryConfigurationException {
  		this.context = ctx;
  		
  		scheduler = Executors.newScheduledThreadPool(Runtime.getRuntime().availableProcessors() * 4);
  		
  		synchronized (MessageCacheTraceRegistry.class) {
  			if(responseCacheManager == null) {
  				InputStream ehcacheConfig= getClass().getResourceAsStream("/ehcache.xml");
  				if(ehcacheConfig == null) {
  					throw new IllegalStateException("Unable to get cache configuration.");
  				}
  				responseCacheManager = CacheManager.create(ehcacheConfig);
  				try {
  					ehcacheConfig.close();
  				} catch (Exception e) {
  					ctx.getLogger().warning("Unable to close cache configuration.  May not affect run");
  				}
  				StringBuilder builder = new StringBuilder();
  				builder.append("Configured Caches:").append('\n');
  				for(String name : responseCacheManager.getCacheNames()) {
  					builder.append('\t').append(name).append('\n');
  				}
  				ctx.getLogger().info(builder.toString());
  				ctx.getLogger().info("Registered MBean Server: " + responseCacheManager.getName());
  				MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();
  				ManagementService.registerMBeans(responseCacheManager, mBeanServer, false, false, false, true);
  				
  			}
  		}
  		
  		Cache cache = responseCacheManager.getCache("JMSDriverDistributedCache");
  		if(cache == null) {
  			throw new RegistryConfigurationException("Unable to get cache");
  		}
  		
  		this.cache = cache;
  	}
  
  	/**
  	 * @see com.sun.faban.driver.transport.asynchronous.TraceRegistry#compileResults(int, java.util.concurrent.TimeUnit)
  	 */
  	public Object compileResults(int time, TimeUnit unit) {
  		// TODO Auto-generated method stub
  		return null;
  	}
  
  	/**
  	 * @see com.sun.faban.driver.transport.asynchronous.TraceRegistry#isComplete(com.sun.faban.driver.transport.asynchronous.Trace)
  	 */
  	public boolean isComplete(Trace<K> trace) {
  		return cache.isKeyInCache(trace.getIdentifier());
  	}
  	
  	/**
  	 * @param payload The message payload
  	 * @return trace A trace object
  	 * @throws DyeingException 
  	 * @see com.sun.faban.driver.transport.asynchronous.TraceRegistry#registerAndDye(java.lang.Object)
  	 */
  	public Trace<K> registerAndDye(K payload) throws DyeingException {
  		try {
  			String correlationId = payload.getJMSCorrelationID();
  			// set the JMSCorrelationID if it's not set.
  			if(correlationId == null) {
  				correlationId = context.getDriverName() + ":" + 
  				context.getAgentId() + ":" + 
  				Long.toString(traceCounter.incrementAndGet());
  				payload.setJMSCorrelationID(correlationId);
  			} 
  			return new MessageTrace(correlationId, payload);
  		} catch (Exception e) {
  			throw new DyeingException(e);
  		}
  	}
  
  	/**
  	 * @see com.sun.faban.driver.transport.asynchronous.TraceRegistry#waitForCompletion(com.sun.faban.driver.transport.asynchronous.Trace, int, java.util.concurrent.TimeUnit)
  	 */
  	@SuppressWarnings("unchecked")
  	public T waitForCompletion(Trace<K> trace, int time, TimeUnit unit) throws InterruptedException, ExecutionException {
  		Future<T> future;
  		T result = null;
  		Element element = cache.get(trace);
  		if(element == null) {
  			future = scheduler.submit(new PollingCacheMonitor(trace, 5));
  			try { 
  				result = future.get(time, unit); // wait!!!
  			} catch (TimeoutException e) {
  				context.getLogger().log(Level.FINE, e.getMessage(), e);	
  			}
  		} else {
  			result = (T)element.getObjectValue(); // No support for generics
  		}
  		return result;
  	}
  
  	/**
  	 * @see java.lang.Object#toString()
  	 */
  	@Override
  	public String toString() {
  		StringBuilder builder = new StringBuilder();
  		builder.append("================== Frequency Statistics ==================").append('\n');
  		builder.append(cacheStatistics.toString()).append('\n');
  		builder.append("================== Frequency Distribution ================").append('\n');
  		builder.append(roundtripFrequence.toString()).append('\n');
  		builder.append("==========================================================").append('\n');
  		return builder.toString();
  	}
  
  	/**
  	 * @see com.sun.faban.driver.transport.asynchronous.TraceRegistry#acknowledge(java.lang.String, java.lang.Object)
  	 */
  	public void acknowledge(String traceId, T response)
  			throws MissingDyeException {
  		if(traceId == null || traceId.length() == 0) {
  			throw new MissingDyeException();
  		}
  		cache.put(new Element(traceId, response));
  		
  	}
  	
  	/** Message Trace **/
  	private final class MessageTrace implements Trace<K> {
  
  		private String id;
  		private K msg;
  		
  		private MessageTrace(String id, K msg) {
  			this.id = id;
  			this.msg = msg;
  		}
  		
  		/**
  		 * @see com.sun.faban.driver.transport.asynchronous.Trace#getIdentifier()
  		 */
  		public String getIdentifier() {
  			return id;
  		}
  
  		/**
  		 * @see com.sun.faban.driver.transport.asynchronous.Trace#getPayload()
  		 */
  		public K getPayload() {
  			return msg;
  		}
  		
  	}
  	
  	/**
  	 * Pooling response class.
  	 * 
  	 * @author Noah Campbell
  	 */
  	final class PollingCacheMonitor implements Callable<T> {
  		private final String correlationId;
  		private int waitTime;
  		
  		/**
  		 * Construct a {@link PollingCacheMonitor}
  		 * 
  		 * @param correlationId
  		 * @param time in milliseconds
  		 */
  		private PollingCacheMonitor(Trace<K> trace, int time) {
  			this.correlationId = trace.getIdentifier();
  			this.waitTime = time;
  		}
  
  		/**
  		 * @see java.util.concurrent.Callable#call()
  		 */
  		@SuppressWarnings("unchecked")
  		public T call() throws Exception {
  			long start = System.currentTimeMillis();
  			while(!cache.isKeyInCache(correlationId)) {
  				Thread.sleep(waitTime);
  			}
  			Element e = cache.get(correlationId);
  			long duration = System.currentTimeMillis() - start;
  			synchronized(PollingCacheMonitor.class) {
  				cacheStatistics.addValue(duration);
  				roundtripFrequence.addValue(duration);
  			}
  			if(e == null) return null;
  			return (T) e.getObjectValue();
  		}
  	}
  
  	
  	private static final Set<Partition> SUPPORTED_PARTITIONS = new HashSet<Partition>();
  	static {
  		Collections.addAll(SUPPORTED_PARTITIONS, Partition.THREAD, Partition.JVM, Partition.HOST);
  	}
  	
  	/**
  	 * @see com.sun.faban.driver.transport.asynchronous.TraceRegistry#getPartitionTolerance()
  	 */
  	public Set<Partition> getPartitionTolerance() {
  		return SUPPORTED_PARTITIONS;
  	}
  
  	/**
  	 * @see com.sun.faban.driver.transport.asynchronous.TraceRegistry#isSafe(com.sun.faban.driver.transport.asynchronous.Partition)
  	 */
  	public boolean isSafe(Partition partition) {
  		return SUPPORTED_PARTITIONS.contains(partition);
  	}
  
  }
  