package org.jboss.cache.util.internals;
import org.jboss.cache.Cache;
import org.jboss.cache.Fqn;
import org.jboss.cache.InvocationContext;
import org.jboss.cache.RPCManager;
import org.jboss.cache.commands.ReplicableCommand;
import org.jboss.cache.commands.remote.ReplicateCommand;
import org.jboss.cache.commands.tx.PrepareCommand;
import org.jboss.cache.factories.ComponentRegistry;
import org.jboss.cache.io.ByteBuffer;
import org.jboss.cache.marshall.AbstractMarshaller;
import org.jboss.cache.marshall.CommandAwareRpcDispatcher;
import org.jboss.cache.marshall.InactiveRegionAwareRpcDispatcher;
import org.jboss.cache.marshall.Marshaller;
import org.jboss.cache.marshall.RegionalizedMethodCall;
import org.jboss.cache.util.TestingUtil;
import org.jgroups.blocks.RpcDispatcher;
import org.jgroups.util.Buffer;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
/**
* Utility class that notifies when certain commands were asynchronously replicated on secondary cache.
* Especially useful for avaoiding Thread.sleep() statements.
* <p/>
* Usage:
* <pre>
* Cache c1, c2; //these being two async caches
* AsyncReplicationListener listener2 = new AsyncReplicationListener(c2);
* listener2.expect(PutKeyValueCommand.class);
* c1.put(fqn, key, value);
* listener2.waitForReplicationToOccur(1000); // -this will block here untill c2 recieves the PutKeyValueCommand command
* </pre>
* Lifecycle - after being used (i.e. waitForReplicationToOccur returns sucessfully) the object returns to the
* non-initialized state and *can* be reused through expect-wait cycle.
* <b>Note</b>: this class might be used aswell for sync caches, e.g. a test could have subclasses which use sync and
* async replication
*
* @author Mircea.Markus@jboss.com
* @since 2.2
*/
public class ReplicationListener
{
private CountDownLatch latch = new CountDownLatch(1);
private Set<Class<? extends ReplicableCommand>> expectedCommands;
/**
* Builds a listener that will observe the given cache for recieving replication commands.
*/
public ReplicationListener(Cache cache)
{
ComponentRegistry componentRegistry = TestingUtil.extractComponentRegistry(cache);
RPCManager rpcManager = componentRegistry.getComponent(RPCManager.class);
CommandAwareRpcDispatcher realDispatcher = (CommandAwareRpcDispatcher) TestingUtil.extractField(rpcManager, "rpcDispatcher");
RpcDispatcher.Marshaller2 realMarshaller = (RpcDispatcher.Marshaller2) realDispatcher.getMarshaller();
RpcDispatcher.Marshaller2 delegate = null;
if (realDispatcher instanceof InactiveRegionAwareRpcDispatcher)
delegate = new RegionMarshallerDelegate((Marshaller) realMarshaller);
else
delegate = new MarshallerDelegate(realMarshaller);
realDispatcher.setMarshaller(delegate);
realDispatcher.setRequestMarshaller(delegate);
realDispatcher.setResponseMarshaller(delegate);
}
private class MarshallerDelegate implements RpcDispatcher.Marshaller2
{
RpcDispatcher.Marshaller2 marshaller;
private MarshallerDelegate(RpcDispatcher.Marshaller2 marshaller)
{
this.marshaller = marshaller;
}
public byte[] objectToByteBuffer(Object obj) throws Exception
{
return marshaller.objectToByteBuffer(obj);
}
public Object objectFromByteBuffer(byte bytes[]) throws Exception
{
Object result = marshaller.objectFromByteBuffer(bytes);
if (result instanceof ReplicateCommand && expectedCommands != null)
{
ReplicateCommand replicateCommand = (ReplicateCommand) result;
return new ReplicateCommandDelegate(replicateCommand);
}
return result;
}
public Buffer objectToBuffer(Object o) throws Exception
{
return marshaller.objectToBuffer(o);
}
public Object objectFromByteBuffer(byte[] bytes, int i, int i1) throws Exception
{
Object result = marshaller.objectFromByteBuffer(bytes, i, i1);
if (result instanceof ReplicateCommand && expectedCommands != null)
{
ReplicateCommand replicateCommand = (ReplicateCommand) result;
return new ReplicateCommandDelegate(replicateCommand);
}
return result;
}
}
/**
* We want the notification to be performed only *after* the remote command is executed.
*/
private class ReplicateCommandDelegate extends ReplicateCommand
{
ReplicateCommand realOne;
private ReplicateCommandDelegate(ReplicateCommand realOne)
{
this.realOne = realOne;
}
@Override
public Object perform(InvocationContext ctx) throws Throwable
{
try
{
return realOne.perform(ctx);
}
finally
{
System.out.println("Processed command: " + realOne);
Iterator<Class<? extends ReplicableCommand>> it = expectedCommands.iterator();
while (it.hasNext())
{
Class<? extends ReplicableCommand> replicableCommandClass = it.next();
if (realOne.containsCommandType(replicableCommandClass))
{
it.remove();
}
else if (realOne.getSingleModification() instanceof PrepareCommand) //explicit transaction
{
PrepareCommand prepareCommand = (PrepareCommand) realOne.getSingleModification();
if (prepareCommand.containsModificationType(replicableCommandClass))
{
it.remove();
}
}
}
if (expectedCommands.isEmpty())
{
latch.countDown();
}
}
}
}
/**
* Needed for region based marshalling.
*/
private class RegionMarshallerDelegate extends AbstractMarshaller
{
private Marshaller realOne;
private RegionMarshallerDelegate(Marshaller realOne)
{
this.realOne = realOne;
}
public void objectToObjectStream(Object obj, ObjectOutputStream out) throws Exception
{
realOne.objectToObjectStream(obj, out);
}
public Object objectFromObjectStream(ObjectInputStream in) throws Exception
{
return realOne.objectFromObjectStream(in);
}
public Object objectFromStream(InputStream is) throws Exception
{
return realOne.objectFromStream(is);
}
public void objectToObjectStream(Object obj, ObjectOutputStream out, Fqn region) throws Exception
{
realOne.objectToObjectStream(obj, out, region);
}
public RegionalizedMethodCall regionalizedMethodCallFromByteBuffer(byte[] buffer) throws Exception
{
RegionalizedMethodCall result = realOne.regionalizedMethodCallFromByteBuffer(buffer);
if (result.command instanceof ReplicateCommand && expectedCommands != null)
{
ReplicateCommand replicateCommand = (ReplicateCommand) result.command;
result.command = new ReplicateCommandDelegate(replicateCommand);
}
return result;
}
public RegionalizedMethodCall regionalizedMethodCallFromObjectStream(ObjectInputStream in) throws Exception
{
return realOne.regionalizedMethodCallFromObjectStream(in);
}
public ByteBuffer objectToBuffer(Object o) throws Exception
{
return realOne.objectToBuffer(o);
}
public Object objectFromByteBuffer(byte[] bytes, int i, int i1) throws Exception
{
return realOne.objectFromByteBuffer(bytes, i, i1);
}
}
/**
* Blocks for the elements specified through {@link #expect(Class[])} invocations to be replicated in this cache.
* if replication does not occur in the give timeout then an exception is being thrown.
*/
public void waitForReplicationToOccur(long timeoutMillis)
{
System.out.println("enter... ReplicationListener.waitForReplicationToOccur");
waitForReplicationToOccur(timeoutMillis, TimeUnit.MILLISECONDS);
System.out.println("exit... ReplicationListener.waitForReplicationToOccur");
}
/**
* Similar to {@link #waitForReplicationToOccur(long)} except that this method provides more flexibility in time units.
*
* @param timeout the maximum time to wait
* @param timeUnit the time unit of the <tt>timeout</tt> argument.
*/
public void waitForReplicationToOccur(long timeout, TimeUnit timeUnit)
{
assert expectedCommands != null : "there are no replication expectations; please use AsyncReplicationListener.expect(...) before calling this method";
try
{
if (!latch.await(timeout, timeUnit))
{
assert false : "waiting for more than " + timeout + " " + timeUnit + " and following commands did not replicate: " + expectedCommands;
}
}
catch (InterruptedException e)
{
throw new IllegalStateException("unexpected", e);
}
finally
{
expectedCommands = null;
latch = new CountDownLatch(1);
}
}
/**
* {@link #waitForReplicationToOccur(long)} will block untill all the commands specified here are being replicated
* to this cache. The method can be called several times with various arguments.
*/
public void expect(Class<? extends ReplicableCommand>... expectedCommands)
{
if (this.expectedCommands == null)
{
this.expectedCommands = new HashSet<Class<? extends ReplicableCommand>>();
}
this.expectedCommands.addAll(Arrays.asList(expectedCommands));
}
/**
* Waits untill first command is replicated.
*/
public void expectAny()
{
expect();
}
}