package org.javaee7.jms.send.receive.mdb;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import javax.annotation.Resource;
import javax.interceptor.AroundInvoke;
import javax.interceptor.InvocationContext;
import javax.transaction.Status;
import javax.transaction.Synchronization;
import javax.transaction.TransactionSynchronizationRegistry;
/**
* Allows test to wait until a method is invoked.
*
* @author Patrik Dudits
*/
public class ReceptionSynchronizer {
private final static Map<Method, CountDownLatch> barrier = new HashMap<>();
@Resource
TransactionSynchronizationRegistry txRegistry;
@AroundInvoke
public Object invoke(final InvocationContext ctx) throws Exception {
boolean transactional = false;
try {
System.out.println("Intercepting " + ctx.getMethod().toGenericString());
transactional = txRegistry != null && txRegistry.getTransactionStatus() != Status.STATUS_NO_TRANSACTION;
if (transactional) {
txRegistry.registerInterposedSynchronization(new Synchronization() {
@Override
public void beforeCompletion() {
}
@Override
public void afterCompletion(int i) {
registerInvocation(ctx.getMethod());
}
});
}
return ctx.proceed();
} finally {
if (!transactional) {
registerInvocation(ctx.getMethod());
}
}
}
void registerInvocation(Method m) {
CountDownLatch latch = null;
synchronized (barrier) {
if (barrier.containsKey(m)) {
latch = barrier.get(m);
}
}
if (latch != null) {
latch.countDown();
}
}
public static void waitFor(Class<?> clazz, String methodName) throws InterruptedException {
Method method = null;
for (Method declaredMethod : clazz.getDeclaredMethods()) {
if (methodName.equals(declaredMethod.getName())) {
if (method == null) {
method = declaredMethod;
} else {
throw new IllegalArgumentException(methodName + " is not unique in " + clazz.getSimpleName());
}
}
}
if (method == null) {
throw new IllegalArgumentException(methodName + " not found in " + clazz.getSimpleName());
}
waitFor(method);
}
private static void waitFor(Method method) throws InterruptedException {
CountDownLatch latch;
synchronized (barrier) {
if (barrier.containsKey(method)) {
latch = barrier.get(method);
if (latch.getCount() == 0) {
throw new IllegalStateException("The invocation already happened");
} else {
throw new IllegalStateException("Sorry, I only serve the first one");
}
} else {
latch = new CountDownLatch(1);
barrier.put(method, latch);
}
}
if (!latch.await(2, TimeUnit.SECONDS)) {
throw new AssertionError("Expected method has not been invoked");
}
}
}